cranelift_entity/
primary.rs

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

cranelift_entity/
primary.rs
