wasmer_types/entity/
secondary_map.rs

1// This file contains code from external sources.
2// Attributions: https://github.com/wasmerio/wasmer/blob/main/docs/ATTRIBUTIONS.md
3
4//! Densely numbered entity references as mapping keys.
5
6use crate::entity::iter::{Iter, IterMut};
7use crate::entity::keys::Keys;
8use crate::entity::EntityRef;
9use crate::lib::std::cmp::min;
10use crate::lib::std::marker::PhantomData;
11use crate::lib::std::ops::{Index, IndexMut};
12use crate::lib::std::slice;
13use crate::lib::std::vec::Vec;
14use rkyv::{Archive, Archived, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
15#[cfg(feature = "enable-serde")]
16use serde::{
17    de::{Deserializer, SeqAccess, Visitor},
18    ser::{SerializeSeq, Serializer},
19    Deserialize, Serialize,
20};
21
22/// A mapping `K -> V` for densely indexed entity references.
23///
24/// The `SecondaryMap` data structure uses the dense index space to implement a map with a vector.
25/// Unlike `PrimaryMap`, an `SecondaryMap` can't be used to allocate entity references. It is used
26/// to associate secondary information with entities.
27///
28/// The map does not track if an entry for a key has been inserted or not. Instead it behaves as if
29/// all keys have a default entry from the beginning.
30#[derive(Debug, Clone, RkyvSerialize, RkyvDeserialize, Archive)]
31#[archive_attr(derive(rkyv::CheckBytes))]
32pub struct SecondaryMap<K, V>
33where
34    K: EntityRef,
35    V: Clone,
36{
37    pub(crate) elems: Vec<V>,
38    pub(crate) default: V,
39    pub(crate) unused: PhantomData<K>,
40}
41
42#[cfg(feature = "artifact-size")]
43impl<K, V> loupe::MemoryUsage for SecondaryMap<K, V>
44where
45    K: EntityRef,
46    V: Clone + loupe::MemoryUsage,
47{
48    fn size_of_val(&self, tracker: &mut dyn loupe::MemoryUsageTracker) -> usize {
49        std::mem::size_of_val(self)
50            + self
51                .elems
52                .iter()
53                .map(|value| value.size_of_val(tracker) - std::mem::size_of_val(value))
54                .sum::<usize>()
55    }
56}
57
58/// Shared `SecondaryMap` implementation for all value types.
59impl<K, V> SecondaryMap<K, V>
60where
61    K: EntityRef,
62    V: Clone,
63{
64    /// Create a new empty map.
65    pub fn new() -> Self
66    where
67        V: Default,
68    {
69        Self {
70            elems: Vec::new(),
71            default: Default::default(),
72            unused: PhantomData,
73        }
74    }
75
76    /// Create a new, empty map with the specified capacity.
77    ///
78    /// The map will be able to hold exactly `capacity` elements without reallocating.
79    pub fn with_capacity(capacity: usize) -> Self
80    where
81        V: Default,
82    {
83        Self {
84            elems: Vec::with_capacity(capacity),
85            default: Default::default(),
86            unused: PhantomData,
87        }
88    }
89
90    /// Create a new empty map with a specified default value.
91    ///
92    /// This constructor does not require V to implement Default.
93    pub fn with_default(default: V) -> Self {
94        Self {
95            elems: Vec::new(),
96            default,
97            unused: PhantomData,
98        }
99    }
100
101    /// Returns the number of elements the map can hold without reallocating.
102    pub fn capacity(&self) -> usize {
103        self.elems.capacity()
104    }
105
106    /// Get the element at `k` if it exists.
107    #[inline(always)]
108    pub fn get(&self, k: K) -> Option<&V> {
109        self.elems.get(k.index())
110    }
111
112    /// Is this map completely empty?
113    #[inline(always)]
114    pub fn is_empty(&self) -> bool {
115        self.elems.is_empty()
116    }
117
118    /// Remove all entries from this map.
119    #[inline(always)]
120    pub fn clear(&mut self) {
121        self.elems.clear()
122    }
123
124    /// Iterate over all the keys and values in this map.
125    pub fn iter(&self) -> Iter<K, V> {
126        Iter::new(self.elems.iter())
127    }
128
129    /// Iterate over all the keys and values in this map, mutable edition.
130    pub fn iter_mut(&mut self) -> IterMut<K, V> {
131        IterMut::new(self.elems.iter_mut())
132    }
133
134    /// Iterate over all the keys in this map.
135    pub fn keys(&self) -> Keys<K> {
136        Keys::with_len(self.elems.len())
137    }
138
139    /// Iterate over all the values in this map.
140    pub fn values(&self) -> slice::Iter<V> {
141        self.elems.iter()
142    }
143
144    /// Iterate over all the values in this map, mutable edition.
145    pub fn values_mut(&mut self) -> slice::IterMut<V> {
146        self.elems.iter_mut()
147    }
148
149    /// Resize the map to have `n` entries by adding default entries as needed.
150    pub fn resize(&mut self, n: usize) {
151        self.elems.resize(n, self.default.clone());
152    }
153}
154
155impl<K, V> ArchivedSecondaryMap<K, V>
156where
157    K: EntityRef,
158    V: Archive + Clone,
159{
160    /// Get the element at `k` if it exists.
161    pub fn get(&self, k: K) -> Option<&V::Archived> {
162        self.elems.get(k.index())
163    }
164
165    /// Iterator over all values in the `ArchivedPrimaryMap`
166    pub fn values(&self) -> slice::Iter<Archived<V>> {
167        self.elems.iter()
168    }
169
170    /// Iterate over all the keys and values in this map.
171    pub fn iter(&self) -> Iter<K, Archived<V>> {
172        Iter::new(self.elems.iter())
173    }
174}
175
176impl<K, V> std::fmt::Debug for ArchivedSecondaryMap<K, V>
177where
178    K: EntityRef + std::fmt::Debug,
179    V: Archive + Clone,
180    V::Archived: std::fmt::Debug,
181{
182    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
183        f.debug_map().entries(self.iter()).finish()
184    }
185}
186
187impl<K, V> Default for SecondaryMap<K, V>
188where
189    K: EntityRef,
190    V: Clone + Default,
191{
192    fn default() -> Self {
193        Self::new()
194    }
195}
196
197/// Immutable indexing into an `SecondaryMap`.
198///
199/// All keys are permitted. Untouched entries have the default value.
200impl<K, V> Index<K> for SecondaryMap<K, V>
201where
202    K: EntityRef,
203    V: Clone,
204{
205    type Output = V;
206
207    #[inline(always)]
208    fn index(&self, k: K) -> &V {
209        self.elems.get(k.index()).unwrap_or(&self.default)
210    }
211}
212
213/// Mutable indexing into an `SecondaryMap`.
214///
215/// The map grows as needed to accommodate new keys.
216impl<K, V> IndexMut<K> for SecondaryMap<K, V>
217where
218    K: EntityRef,
219    V: Clone,
220{
221    #[inline(always)]
222    fn index_mut(&mut self, k: K) -> &mut V {
223        let i = k.index();
224        if i >= self.elems.len() {
225            self.elems.resize(i + 1, self.default.clone());
226        }
227        &mut self.elems[i]
228    }
229}
230
231impl<K, V> PartialEq for SecondaryMap<K, V>
232where
233    K: EntityRef,
234    V: Clone + PartialEq,
235{
236    fn eq(&self, other: &Self) -> bool {
237        let min_size = min(self.elems.len(), other.elems.len());
238        self.default == other.default
239            && self.elems[..min_size] == other.elems[..min_size]
240            && self.elems[min_size..].iter().all(|e| *e == self.default)
241            && other.elems[min_size..].iter().all(|e| *e == other.default)
242    }
243}
244
245impl<K, V> Eq for SecondaryMap<K, V>
246where
247    K: EntityRef,
248    V: Clone + PartialEq + Eq,
249{
250}
251
252#[cfg(feature = "enable-serde")]
253impl<K, V> Serialize for SecondaryMap<K, V>
254where
255    K: EntityRef,
256    V: Clone + PartialEq + Serialize,
257{
258    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
259    where
260        S: Serializer,
261    {
262        // TODO: bincode encodes option as "byte for Some/None" and then optionally the content
263        // TODO: we can actually optimize it by encoding manually bitmask, then elements
264        let mut elems_cnt = self.elems.len();
265        while elems_cnt > 0 && self.elems[elems_cnt - 1] == self.default {
266            elems_cnt -= 1;
267        }
268        let mut seq = serializer.serialize_seq(Some(1 + elems_cnt))?;
269        seq.serialize_element(&Some(self.default.clone()))?;
270        for e in self.elems.iter().take(elems_cnt) {
271            let some_e = Some(e);
272            seq.serialize_element(if *e == self.default { &None } else { &some_e })?;
273        }
274        seq.end()
275    }
276}
277
278#[cfg(feature = "enable-serde")]
279impl<'de, K, V> Deserialize<'de> for SecondaryMap<K, V>
280where
281    K: EntityRef,
282    V: Clone + Deserialize<'de>,
283{
284    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
285    where
286        D: Deserializer<'de>,
287    {
288        use crate::lib::std::fmt;
289
290        struct SecondaryMapVisitor<K, V> {
291            unused: PhantomData<fn(K) -> V>,
292        }
293
294        impl<'de, K, V> Visitor<'de> for SecondaryMapVisitor<K, V>
295        where
296            K: EntityRef,
297            V: Clone + Deserialize<'de>,
298        {
299            type Value = SecondaryMap<K, V>;
300
301            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
302                formatter.write_str("struct SecondaryMap")
303            }
304
305            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
306            where
307                A: SeqAccess<'de>,
308            {
309                match seq.next_element()? {
310                    Some(Some(default_val)) => {
311                        let default_val: V = default_val; // compiler can't infer the type
312                        let mut m = SecondaryMap::with_default(default_val.clone());
313                        let mut idx = 0;
314                        while let Some(val) = seq.next_element()? {
315                            let val: Option<_> = val; // compiler can't infer the type
316                            m[K::new(idx)] = val.unwrap_or_else(|| default_val.clone());
317                            idx += 1;
318                        }
319                        Ok(m)
320                    }
321                    _ => Err(serde::de::Error::custom("Default value required")),
322                }
323            }
324        }
325
326        deserializer.deserialize_seq(SecondaryMapVisitor {
327            unused: PhantomData {},
328        })
329    }
330}
331
332#[cfg(test)]
333mod tests {
334    use super::*;
335
336    // `EntityRef` impl for testing.
337    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
338    struct E(u32);
339
340    impl EntityRef for E {
341        fn new(i: usize) -> Self {
342            Self(i as u32)
343        }
344        fn index(self) -> usize {
345            self.0 as usize
346        }
347    }
348
349    #[test]
350    fn basic() {
351        let r0 = E(0);
352        let r1 = E(1);
353        let r2 = E(2);
354        let mut m = SecondaryMap::new();
355
356        let v: Vec<E> = m.keys().collect();
357        assert_eq!(v, []);
358
359        m[r2] = 3;
360        m[r1] = 5;
361
362        assert_eq!(m[r1], 5);
363        assert_eq!(m[r2], 3);
364
365        let v: Vec<E> = m.keys().collect();
366        assert_eq!(v, [r0, r1, r2]);
367
368        let shared = &m;
369        assert_eq!(shared[r0], 0);
370        assert_eq!(shared[r1], 5);
371        assert_eq!(shared[r2], 3);
372    }
373}
wasmer_types/entity/secondary_map.rs

wasmer_types/entity/
secondary_map.rs
