wasmer_types/entity/
primary_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.
5use crate::entity::boxed_slice::BoxedSlice;
6use crate::entity::iter::{IntoIter, Iter, IterMut};
7use crate::entity::keys::Keys;
8use crate::entity::EntityRef;
9use crate::lib::std::boxed::Box;
10use crate::lib::std::iter::FromIterator;
11use crate::lib::std::marker::PhantomData;
12use crate::lib::std::ops::{Index, IndexMut};
13use crate::lib::std::slice;
14use crate::lib::std::vec::Vec;
15use rkyv::{Archive, Archived, Deserialize as RkyvDeserialize, Serialize as RkyvSerialize};
16#[cfg(feature = "enable-serde")]
17use serde::{Deserialize, Serialize};
18
19/// A primary mapping `K -> V` allocating dense entity references.
20///
21/// The `PrimaryMap` data structure uses the dense index space to implement a map with a vector.
22///
23/// A primary map contains the main definition of an entity, and it can be used to allocate new
24/// entity references with the `push` method.
25///
26/// There should only be a single `PrimaryMap` instance for a given `EntityRef` type, otherwise
27/// conflicting references will be created. Using unknown keys for indexing will cause a panic.
28///
29/// Note that `PrimaryMap` doesn't implement `Deref` or `DerefMut`, which would allow
30/// `&PrimaryMap<K, V>` to convert to `&[V]`. One of the main advantages of `PrimaryMap` is
31/// that it only allows indexing with the distinct `EntityRef` key type, so converting to a
32/// plain slice would make it easier to use incorrectly. To make a slice of a `PrimaryMap`, use
33/// `into_boxed_slice`.
34#[derive(Debug, Clone, Hash, PartialEq, Eq)]
35#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
36#[derive(RkyvSerialize, RkyvDeserialize, Archive)]
37#[archive_attr(derive(rkyv::CheckBytes))]
38pub struct PrimaryMap<K, V>
39where
40    K: EntityRef,
41{
42    pub(crate) elems: Vec<V>,
43    pub(crate) unused: PhantomData<K>,
44}
45
46#[cfg(feature = "artifact-size")]
47impl<K, V> loupe::MemoryUsage for PrimaryMap<K, V>
48where
49    K: EntityRef,
50    V: loupe::MemoryUsage,
51{
52    fn size_of_val(&self, tracker: &mut dyn loupe::MemoryUsageTracker) -> usize {
53        std::mem::size_of_val(self)
54            + self
55                .elems
56                .iter()
57                .map(|value| value.size_of_val(tracker) - std::mem::size_of_val(value))
58                .sum::<usize>()
59    }
60}
61
62impl<K, V> PrimaryMap<K, V>
63where
64    K: EntityRef,
65{
66    /// Create a new empty map.
67    pub fn new() -> Self {
68        Self {
69            elems: Vec::new(),
70            unused: PhantomData,
71        }
72    }
73
74    /// Create a new empty map with the given capacity.
75    pub fn with_capacity(capacity: usize) -> Self {
76        Self {
77            elems: Vec::with_capacity(capacity),
78            unused: PhantomData,
79        }
80    }
81
82    /// Check if `k` is a valid key in the map.
83    pub fn is_valid(&self, k: K) -> bool {
84        k.index() < self.elems.len()
85    }
86
87    /// Get the element at `k` if it exists.
88    pub fn get(&self, k: K) -> Option<&V> {
89        self.elems.get(k.index())
90    }
91
92    /// Get the element at `k` if it exists, mutable version.
93    pub fn get_mut(&mut self, k: K) -> Option<&mut V> {
94        self.elems.get_mut(k.index())
95    }
96
97    /// Is this map completely empty?
98    pub fn is_empty(&self) -> bool {
99        self.elems.is_empty()
100    }
101
102    /// Get the total number of entity references created.
103    pub fn len(&self) -> usize {
104        self.elems.len()
105    }
106
107    /// Iterate over all the keys in this map.
108    pub fn keys(&self) -> Keys<K> {
109        Keys::with_len(self.elems.len())
110    }
111
112    /// Iterate over all the values in this map.
113    pub fn values(&self) -> slice::Iter<V> {
114        self.elems.iter()
115    }
116
117    /// Iterate over all the values in this map, mutable edition.
118    pub fn values_mut(&mut self) -> slice::IterMut<V> {
119        self.elems.iter_mut()
120    }
121
122    /// Iterate over all the keys and values in this map.
123    pub fn iter(&self) -> Iter<K, V> {
124        Iter::new(self.elems.iter())
125    }
126
127    /// Iterate over all the keys and values in this map, mutable edition.
128    pub fn iter_mut(&mut self) -> IterMut<K, V> {
129        IterMut::new(self.elems.iter_mut())
130    }
131
132    /// Remove all entries from this map.
133    pub fn clear(&mut self) {
134        self.elems.clear()
135    }
136
137    /// Get the key that will be assigned to the next pushed value.
138    pub fn next_key(&self) -> K {
139        K::new(self.elems.len())
140    }
141
142    /// Append `v` to the mapping, assigning a new key which is returned.
143    pub fn push(&mut self, v: V) -> K {
144        let k = self.next_key();
145        self.elems.push(v);
146        k
147    }
148
149    /// Returns the last element that was inserted in the map.
150    pub fn last(&self) -> Option<&V> {
151        self.elems.last()
152    }
153
154    /// Reserves capacity for at least `additional` more elements to be inserted.
155    pub fn reserve(&mut self, additional: usize) {
156        self.elems.reserve(additional)
157    }
158
159    /// Reserves the minimum capacity for exactly `additional` more elements to be inserted.
160    pub fn reserve_exact(&mut self, additional: usize) {
161        self.elems.reserve_exact(additional)
162    }
163
164    /// Shrinks the capacity of the `PrimaryMap` as much as possible.
165    pub fn shrink_to_fit(&mut self) {
166        self.elems.shrink_to_fit()
167    }
168
169    /// Consumes this `PrimaryMap` and produces a `BoxedSlice`.
170    pub fn into_boxed_slice(self) -> BoxedSlice<K, V> {
171        unsafe { BoxedSlice::<K, V>::from_raw(Box::<[V]>::into_raw(self.elems.into_boxed_slice())) }
172    }
173}
174
175impl<K, V> ArchivedPrimaryMap<K, V>
176where
177    K: EntityRef,
178    V: Archive,
179{
180    /// Get the element at `k` if it exists.
181    pub fn get(&self, k: K) -> Option<&V::Archived> {
182        self.elems.get(k.index())
183    }
184}
185
186impl<K, V> std::fmt::Debug for ArchivedPrimaryMap<K, V>
187where
188    K: EntityRef + std::fmt::Debug,
189    V: Archive,
190    V::Archived: std::fmt::Debug,
191{
192    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
193        f.debug_map().entries(self.iter()).finish()
194    }
195}
196
197impl<K, V> Default for PrimaryMap<K, V>
198where
199    K: EntityRef,
200{
201    fn default() -> Self {
202        Self::new()
203    }
204}
205
206/// Immutable indexing into an `PrimaryMap`.
207/// The indexed value must be in the map.
208impl<K, V> Index<K> for PrimaryMap<K, V>
209where
210    K: EntityRef,
211{
212    type Output = V;
213
214    fn index(&self, k: K) -> &V {
215        &self.elems[k.index()]
216    }
217}
218
219/// Mutable indexing into an `PrimaryMap`.
220impl<K, V> IndexMut<K> for PrimaryMap<K, V>
221where
222    K: EntityRef,
223{
224    fn index_mut(&mut self, k: K) -> &mut V {
225        &mut self.elems[k.index()]
226    }
227}
228
229impl<K, V> IntoIterator for PrimaryMap<K, V>
230where
231    K: EntityRef,
232{
233    type Item = (K, V);
234    type IntoIter = IntoIter<K, V>;
235
236    fn into_iter(self) -> Self::IntoIter {
237        IntoIter::new(self.elems.into_iter())
238    }
239}
240
241impl<'a, K, V> IntoIterator for &'a PrimaryMap<K, V>
242where
243    K: EntityRef,
244{
245    type Item = (K, &'a V);
246    type IntoIter = Iter<'a, K, V>;
247
248    fn into_iter(self) -> Self::IntoIter {
249        Iter::new(self.elems.iter())
250    }
251}
252
253impl<'a, K, V> IntoIterator for &'a mut PrimaryMap<K, V>
254where
255    K: EntityRef,
256{
257    type Item = (K, &'a mut V);
258    type IntoIter = IterMut<'a, K, V>;
259
260    fn into_iter(self) -> Self::IntoIter {
261        IterMut::new(self.elems.iter_mut())
262    }
263}
264
265impl<K, V> FromIterator<V> for PrimaryMap<K, V>
266where
267    K: EntityRef,
268{
269    fn from_iter<T>(iter: T) -> Self
270    where
271        T: IntoIterator<Item = V>,
272    {
273        Self {
274            elems: Vec::from_iter(iter),
275            unused: PhantomData,
276        }
277    }
278}
279
280impl<K, V> ArchivedPrimaryMap<K, V>
281where
282    K: EntityRef,
283    V: Archive,
284    V::Archived: std::fmt::Debug,
285{
286    /// Iterator over all values in the `ArchivedPrimaryMap`
287    pub fn values(&self) -> slice::Iter<Archived<V>> {
288        self.elems.iter()
289    }
290
291    /// Iterate over all the keys and values in this map.
292    pub fn iter(&self) -> Iter<K, Archived<V>> {
293        Iter::new(self.elems.iter())
294    }
295}
296
297/// Immutable indexing into an `ArchivedPrimaryMap`.
298/// The indexed value must be in the map.
299impl<K, V> Index<K> for ArchivedPrimaryMap<K, V>
300where
301    K: EntityRef,
302    V: Archive,
303    V::Archived: std::fmt::Debug,
304{
305    type Output = Archived<V>;
306
307    fn index(&self, k: K) -> &Self::Output {
308        &self.elems[k.index()]
309    }
310}
311
312#[cfg(test)]
313mod tests {
314    use super::*;
315
316    // `EntityRef` impl for testing.
317    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
318    struct E(u32);
319
320    impl EntityRef for E {
321        fn new(i: usize) -> Self {
322            Self(i as u32)
323        }
324        fn index(self) -> usize {
325            self.0 as usize
326        }
327    }
328
329    #[test]
330    fn basic() {
331        let r0 = E(0);
332        let r1 = E(1);
333        let m = PrimaryMap::<E, isize>::new();
334
335        let v: Vec<E> = m.keys().collect();
336        assert_eq!(v, []);
337
338        assert!(!m.is_valid(r0));
339        assert!(!m.is_valid(r1));
340    }
341
342    #[test]
343    fn push() {
344        let mut m = PrimaryMap::new();
345        let k0: E = m.push(12);
346        let k1 = m.push(33);
347
348        assert_eq!(m[k0], 12);
349        assert_eq!(m[k1], 33);
350
351        let v: Vec<E> = m.keys().collect();
352        assert_eq!(v, [k0, k1]);
353    }
354
355    #[test]
356    fn iter() {
357        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
358        m.push(12);
359        m.push(33);
360
361        let mut i = 0;
362        for (key, value) in &m {
363            assert_eq!(key.index(), i);
364            match i {
365                0 => assert_eq!(*value, 12),
366                1 => assert_eq!(*value, 33),
367                _ => panic!(),
368            }
369            i += 1;
370        }
371        i = 0;
372        for (key_mut, value_mut) in m.iter_mut() {
373            assert_eq!(key_mut.index(), i);
374            match i {
375                0 => assert_eq!(*value_mut, 12),
376                1 => assert_eq!(*value_mut, 33),
377                _ => panic!(),
378            }
379            i += 1;
380        }
381    }
382
383    #[test]
384    fn iter_rev() {
385        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
386        m.push(12);
387        m.push(33);
388
389        let mut i = 2;
390        for (key, value) in m.iter().rev() {
391            i -= 1;
392            assert_eq!(key.index(), i);
393            match i {
394                0 => assert_eq!(*value, 12),
395                1 => assert_eq!(*value, 33),
396                _ => panic!(),
397            }
398        }
399
400        i = 2;
401        for (key, value) in m.iter_mut().rev() {
402            i -= 1;
403            assert_eq!(key.index(), i);
404            match i {
405                0 => assert_eq!(*value, 12),
406                1 => assert_eq!(*value, 33),
407                _ => panic!(),
408            }
409        }
410    }
411    #[test]
412    fn keys() {
413        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
414        m.push(12);
415        m.push(33);
416
417        for (i, key) in m.keys().enumerate() {
418            assert_eq!(key.index(), i);
419        }
420    }
421
422    #[test]
423    fn keys_rev() {
424        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
425        m.push(12);
426        m.push(33);
427
428        let mut i = 2;
429        for key in m.keys().rev() {
430            i -= 1;
431            assert_eq!(key.index(), i);
432        }
433    }
434
435    #[test]
436    fn values() {
437        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
438        m.push(12);
439        m.push(33);
440
441        let mut i = 0;
442        for value in m.values() {
443            match i {
444                0 => assert_eq!(*value, 12),
445                1 => assert_eq!(*value, 33),
446                _ => panic!(),
447            }
448            i += 1;
449        }
450        i = 0;
451        for value_mut in m.values_mut() {
452            match i {
453                0 => assert_eq!(*value_mut, 12),
454                1 => assert_eq!(*value_mut, 33),
455                _ => panic!(),
456            }
457            i += 1;
458        }
459    }
460
461    #[test]
462    fn values_rev() {
463        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
464        m.push(12);
465        m.push(33);
466
467        let mut i = 2;
468        for value in m.values().rev() {
469            i -= 1;
470            match i {
471                0 => assert_eq!(*value, 12),
472                1 => assert_eq!(*value, 33),
473                _ => panic!(),
474            }
475        }
476        i = 2;
477        for value_mut in m.values_mut().rev() {
478            i -= 1;
479            match i {
480                0 => assert_eq!(*value_mut, 12),
481                1 => assert_eq!(*value_mut, 33),
482                _ => panic!(),
483            }
484        }
485    }
486
487    #[test]
488    fn from_iter() {
489        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
490        m.push(12);
491        m.push(33);
492
493        let n = m.values().collect::<PrimaryMap<E, _>>();
494        assert!(m.len() == n.len());
495        for (me, ne) in m.values().zip(n.values()) {
496            assert!(*me == **ne);
497        }
498    }
499}
wasmer_types/entity/primary_map.rs

wasmer_types/entity/
primary_map.rs
