Skip to main content

wasmer_types/entity/
primary_map.rs

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