wasmtime_environ/collections/
secondary_map.rs

1use crate::{error::OutOfMemory, prelude::*};
2use core::{fmt, ops::Index};
3use cranelift_entity::{EntityRef, SecondaryMap as Inner};
4
5/// Like [`cranelift_entity::SecondaryMap`] but all allocation is fallible.
6pub struct TrySecondaryMap<K, V>
7where
8    K: EntityRef,
9    V: Clone,
10{
11    inner: Inner<K, V>,
12}
13
14impl<K, V> fmt::Debug for TrySecondaryMap<K, V>
15where
16    K: EntityRef + fmt::Debug,
17    V: fmt::Debug + Clone,
18{
19    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
20        fmt::Debug::fmt(&self.inner, f)
21    }
22}
23
24impl<K, V> TrySecondaryMap<K, V>
25where
26    K: EntityRef,
27    V: Clone,
28{
29    /// Same as [`cranelift_entity::SecondaryMap::new`].
30    ///
31    /// XXX: `Clone::clone(&<V as Default>::default())` should never allocate.
32    pub fn new() -> Self
33    where
34        V: Default,
35    {
36        Self {
37            inner: Inner::new(),
38        }
39    }
40
41    /// Same as [`cranelift_entity::SecondaryMap::try_with_capacity`].
42    ///
43    /// XXX: `Clone::clone(&<V as Default>::default())` should never allocate.
44    pub fn with_capacity(capacity: usize) -> Result<Self, OutOfMemory>
45    where
46        V: Default,
47    {
48        Ok(Self {
49            inner: Inner::try_with_capacity(capacity)?,
50        })
51    }
52
53    /// Same as [`cranelift_entity::SecondaryMap::with_default`].
54    ///
55    /// XXX: `Clone::clone(&default)` should never allocate.
56    pub fn with_default(default: V) -> Self {
57        Self {
58            inner: Inner::with_default(default),
59        }
60    }
61
62    /// Same as [`cranelift_entity::SecondaryMap::capacity`].
63    pub fn capacity(&self) -> usize {
64        self.inner.capacity()
65    }
66
67    /// Same as [`cranelift_entity::SecondaryMap::get`].
68    pub fn get(&self, k: K) -> Option<&V> {
69        self.inner.get(k)
70    }
71
72    /// Same as [`cranelift_entity::SecondaryMap::get_mut`].
73    pub fn get_mut(&mut self, k: K) -> Option<&mut V> {
74        self.inner.get_mut(k)
75    }
76
77    /// Same as [`cranelift_entity::SecondaryMap::try_insert`].
78    pub fn insert(&mut self, k: K, v: V) -> Result<Option<V>, OutOfMemory> {
79        self.inner.try_insert(k, v)
80    }
81
82    /// Same as [`cranelift_entity::SecondaryMap::remove`].
83    pub fn remove(&mut self, k: K) -> Option<V> {
84        self.inner.remove(k)
85    }
86
87    /// Same as [`cranelift_entity::SecondaryMap::is_empty`].
88    pub fn is_empty(&self) -> bool {
89        self.inner.is_empty()
90    }
91
92    /// Same as [`cranelift_entity::SecondaryMap::clear`].
93    pub fn clear(&mut self) {
94        self.inner.clear()
95    }
96
97    /// Same as [`cranelift_entity::SecondaryMap::iter`].
98    pub fn iter(&self) -> cranelift_entity::Iter<'_, K, V> {
99        self.inner.iter()
100    }
101
102    /// Same as [`cranelift_entity::SecondaryMap::iter_mut`].
103    pub fn iter_mut(&mut self) -> cranelift_entity::IterMut<'_, K, V> {
104        self.inner.iter_mut()
105    }
106
107    /// Same as [`cranelift_entity::SecondaryMap::keys`].
108    pub fn keys(&self) -> cranelift_entity::Keys<K> {
109        self.inner.keys()
110    }
111
112    /// Same as [`cranelift_entity::SecondaryMap::values`].
113    pub fn values(&self) -> core::slice::Iter<'_, V> {
114        self.inner.values()
115    }
116
117    /// Same as [`cranelift_entity::SecondaryMap::values_mut`].
118    pub fn values_mut(&mut self) -> core::slice::IterMut<'_, V> {
119        self.inner.values_mut()
120    }
121
122    /// Resize the map to have `n` entries by adding default entries as needed.
123    pub fn resize(&mut self, n: usize) -> Result<(), OutOfMemory> {
124        self.inner.try_resize(n)
125    }
126}
127
128impl<K, V> Default for TrySecondaryMap<K, V>
129where
130    K: EntityRef,
131    V: Clone + Default,
132{
133    fn default() -> TrySecondaryMap<K, V> {
134        TrySecondaryMap::new()
135    }
136}
137
138// NB: no `IndexMut` implementation because it requires allocation but the trait
139// doesn't allow for fallibility.
140impl<K, V> Index<K> for TrySecondaryMap<K, V>
141where
142    K: EntityRef,
143    V: Clone,
144{
145    type Output = V;
146
147    fn index(&self, k: K) -> &V {
148        &self.inner[k]
149    }
150}
151
152impl<K, V> From<TryVec<V>> for TrySecondaryMap<K, V>
153where
154    K: EntityRef,
155    V: Clone + Default,
156{
157    fn from(values: TryVec<V>) -> Self {
158        let values: alloc::vec::Vec<V> = values.into();
159        let inner = Inner::from(values);
160        Self { inner }
161    }
162}
163
164impl<K, V> serde::ser::Serialize for TrySecondaryMap<K, V>
165where
166    K: EntityRef,
167    V: Clone + PartialEq + serde::ser::Serialize,
168{
169    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
170    where
171        S: serde::Serializer,
172    {
173        use serde::ser::SerializeSeq as _;
174
175        // Ignore any trailing default values.
176        let mut len = self.inner.as_values_slice().len();
177        while len > 0 && &self[K::new(len - 1)] == self.inner.default_value() {
178            len -= 1;
179        }
180
181        // Plus one for the default value.
182        let mut seq = serializer.serialize_seq(Some(len + 1))?;
183
184        // Always serialize the default value first.
185        seq.serialize_element(self.inner.default_value())?;
186
187        for elem in self.values().take(len) {
188            let elem = if elem == self.inner.default_value() {
189                None
190            } else {
191                Some(elem)
192            };
193            seq.serialize_element(&elem)?;
194        }
195
196        seq.end()
197    }
198}
199
200impl<'de, K, V> serde::de::Deserialize<'de> for TrySecondaryMap<K, V>
201where
202    K: EntityRef,
203    V: Clone + serde::de::Deserialize<'de>,
204{
205    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
206    where
207        D: serde::Deserializer<'de>,
208    {
209        struct Visitor<K, V>(core::marker::PhantomData<fn() -> TrySecondaryMap<K, V>>)
210        where
211            K: EntityRef,
212            V: Clone;
213
214        impl<'de, K, V> serde::de::Visitor<'de> for Visitor<K, V>
215        where
216            K: EntityRef,
217            V: Clone + serde::de::Deserialize<'de>,
218        {
219            type Value = TrySecondaryMap<K, V>;
220
221            fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
222                f.write_str("struct SecondaryMap")
223            }
224
225            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
226            where
227                A: serde::de::SeqAccess<'de>,
228            {
229                // Minus one to account for the default element, which is always
230                // the first in the sequence.
231                let size_hint = seq.size_hint().and_then(|n| n.checked_sub(1));
232
233                let Some(default) = seq.next_element::<V>()? else {
234                    return Err(serde::de::Error::custom("Default value required"));
235                };
236
237                let mut map = TrySecondaryMap::<K, V>::with_default(default.clone());
238
239                if let Some(n) = size_hint {
240                    map.resize(n).map_err(|oom| serde::de::Error::custom(oom))?;
241                }
242
243                let mut idx = 0;
244                while let Some(val) = seq.next_element::<Option<V>>()? {
245                    let key = K::new(idx);
246                    let val = match val {
247                        None => default.clone(),
248                        Some(val) => val,
249                    };
250
251                    map.insert(key, val)
252                        .map_err(|oom| serde::de::Error::custom(oom))?;
253
254                    idx += 1;
255                }
256
257                Ok(map)
258            }
259        }
260
261        deserializer.deserialize_seq(Visitor(core::marker::PhantomData))
262    }
263}
264
265#[cfg(test)]
266mod tests {
267    use super::*;
268    use crate::error::Result;
269    use alloc::vec;
270    use alloc::vec::Vec;
271
272    #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
273    struct K(u32);
274    crate::entity_impl!(K);
275
276    fn k(i: usize) -> K {
277        K::new(i)
278    }
279
280    #[test]
281    fn serialize_deserialize() -> Result<()> {
282        let mut map = TrySecondaryMap::<K, u32>::with_default(99);
283        map.insert(k(0), 33)?;
284        map.insert(k(1), 44)?;
285        map.insert(k(2), 55)?;
286        map.insert(k(3), 99)?;
287        map.insert(k(4), 99)?;
288
289        let bytes = postcard::to_allocvec(&map)?;
290        let map2: TrySecondaryMap<K, u32> = postcard::from_bytes(&bytes)?;
291
292        for i in 0..10 {
293            assert_eq!(map[k(i)], map2[k(i)]);
294        }
295
296        // Trailing default entries were omitted from the serialization.
297        assert_eq!(map2.keys().collect::<Vec<_>>(), vec![k(0), k(1), k(2)]);
298
299        Ok(())
300    }
301}
wasmtime_environ/collections/secondary_map.rs

wasmtime_environ/collections/
secondary_map.rs
