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

wasmer_types/entity/
primary_map.rs
