cranelift_entity/
primary.rs

1//! Densely numbered entity references as mapping keys.
2use crate::EntityRef;
3use crate::boxed_slice::BoxedSlice;
4use crate::iter::{IntoIter, Iter, IterMut};
5use crate::keys::Keys;
6use alloc::boxed::Box;
7use alloc::vec::Vec;
8use core::fmt;
9use core::marker::PhantomData;
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(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    pub fn get_disjoint_mut<const N: usize>(
166        &mut self,
167        indices: [K; N],
168    ) -> Result<[&mut V; N], slice::GetDisjointMutError> {
169        self.elems.get_disjoint_mut(indices.map(|k| k.index()))
170    }
171
172    /// Performs a binary search on the values with a key extraction function.
173    ///
174    /// Assumes that the values are sorted by the key extracted by the function.
175    ///
176    /// If the value is found then `Ok(K)` is returned, containing the entity key
177    /// of the matching value.
178    ///
179    /// If there are multiple matches, then any one of the matches could be returned.
180    ///
181    /// If the value is not found then Err(K) is returned, containing the entity key
182    /// where a matching element could be inserted while maintaining sorted order.
183    pub fn binary_search_values_by_key<'a, B, F>(&'a self, b: &B, f: F) -> Result<K, K>
184    where
185        F: FnMut(&'a V) -> B,
186        B: Ord,
187    {
188        self.elems
189            .binary_search_by_key(b, f)
190            .map(|i| K::new(i))
191            .map_err(|i| K::new(i))
192    }
193
194    /// Analog of `get_raw` except that a raw pointer is returned rather than a
195    /// mutable reference.
196    ///
197    /// The default accessors of items in [`PrimaryMap`] will invalidate all
198    /// previous borrows obtained from the map according to miri. This function
199    /// can be used to acquire a pointer and then subsequently acquire a second
200    /// pointer later on without invalidating the first one. In other words
201    /// this is only here to help borrow two elements simultaneously with miri.
202    pub fn get_raw_mut(&mut self, k: K) -> Option<*mut V> {
203        if k.index() < self.elems.len() {
204            // SAFETY: the `add` function requires that the index is in-bounds
205            // with respect to the allocation which is satisfied here due to
206            // the bounds-check above.
207            unsafe { Some(self.elems.as_mut_ptr().add(k.index())) }
208        } else {
209            None
210        }
211    }
212}
213
214impl<K, V> Default for PrimaryMap<K, V>
215where
216    K: EntityRef,
217{
218    fn default() -> PrimaryMap<K, V> {
219        PrimaryMap::new()
220    }
221}
222
223/// Immutable indexing into an `PrimaryMap`.
224/// The indexed value must be in the map.
225impl<K, V> Index<K> for PrimaryMap<K, V>
226where
227    K: EntityRef,
228{
229    type Output = V;
230
231    fn index(&self, k: K) -> &V {
232        &self.elems[k.index()]
233    }
234}
235
236/// Mutable indexing into an `PrimaryMap`.
237impl<K, V> IndexMut<K> for PrimaryMap<K, V>
238where
239    K: EntityRef,
240{
241    fn index_mut(&mut self, k: K) -> &mut V {
242        &mut self.elems[k.index()]
243    }
244}
245
246impl<K, V> IntoIterator for PrimaryMap<K, V>
247where
248    K: EntityRef,
249{
250    type Item = (K, V);
251    type IntoIter = IntoIter<K, V>;
252
253    fn into_iter(self) -> Self::IntoIter {
254        IntoIter::new(self.elems.into_iter())
255    }
256}
257
258impl<'a, K, V> IntoIterator for &'a PrimaryMap<K, V>
259where
260    K: EntityRef,
261{
262    type Item = (K, &'a V);
263    type IntoIter = Iter<'a, K, V>;
264
265    fn into_iter(self) -> Self::IntoIter {
266        Iter::new(self.elems.iter())
267    }
268}
269
270impl<'a, K, V> IntoIterator for &'a mut PrimaryMap<K, V>
271where
272    K: EntityRef,
273{
274    type Item = (K, &'a mut V);
275    type IntoIter = IterMut<'a, K, V>;
276
277    fn into_iter(self) -> Self::IntoIter {
278        IterMut::new(self.elems.iter_mut())
279    }
280}
281
282impl<K, V> FromIterator<V> for PrimaryMap<K, V>
283where
284    K: EntityRef,
285{
286    fn from_iter<T>(iter: T) -> Self
287    where
288        T: IntoIterator<Item = V>,
289    {
290        Self {
291            elems: Vec::from_iter(iter),
292            unused: PhantomData,
293        }
294    }
295}
296
297impl<K, V> From<Vec<V>> for PrimaryMap<K, V>
298where
299    K: EntityRef,
300{
301    fn from(elems: Vec<V>) -> Self {
302        Self {
303            elems,
304            unused: PhantomData,
305        }
306    }
307}
308
309impl<K: EntityRef + fmt::Debug, V: fmt::Debug> fmt::Debug for PrimaryMap<K, V> {
310    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
311        let mut struct_ = f.debug_struct("PrimaryMap");
312        for (k, v) in self {
313            struct_.field(&alloc::format!("{k:?}"), v);
314        }
315        struct_.finish()
316    }
317}
318
319#[cfg(test)]
320mod tests {
321    use super::*;
322
323    // `EntityRef` impl for testing.
324    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
325    struct E(u32);
326
327    impl EntityRef for E {
328        fn new(i: usize) -> Self {
329            E(i as u32)
330        }
331        fn index(self) -> usize {
332            self.0 as usize
333        }
334    }
335
336    #[test]
337    fn basic() {
338        let r0 = E(0);
339        let r1 = E(1);
340        let m = PrimaryMap::<E, isize>::new();
341
342        let v: Vec<E> = m.keys().collect();
343        assert_eq!(v, []);
344
345        assert!(!m.is_valid(r0));
346        assert!(!m.is_valid(r1));
347    }
348
349    #[test]
350    fn push() {
351        let mut m = PrimaryMap::new();
352        let k0: E = m.push(12);
353        let k1 = m.push(33);
354
355        assert_eq!(m[k0], 12);
356        assert_eq!(m[k1], 33);
357
358        let v: Vec<E> = m.keys().collect();
359        assert_eq!(v, [k0, k1]);
360    }
361
362    #[test]
363    fn iter() {
364        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
365        m.push(12);
366        m.push(33);
367
368        let mut i = 0;
369        for (key, value) in &m {
370            assert_eq!(key.index(), i);
371            match i {
372                0 => assert_eq!(*value, 12),
373                1 => assert_eq!(*value, 33),
374                _ => panic!(),
375            }
376            i += 1;
377        }
378        i = 0;
379        for (key_mut, value_mut) in m.iter_mut() {
380            assert_eq!(key_mut.index(), i);
381            match i {
382                0 => assert_eq!(*value_mut, 12),
383                1 => assert_eq!(*value_mut, 33),
384                _ => panic!(),
385            }
386            i += 1;
387        }
388    }
389
390    #[test]
391    fn iter_rev() {
392        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
393        m.push(12);
394        m.push(33);
395
396        let mut i = 2;
397        for (key, value) in m.iter().rev() {
398            i -= 1;
399            assert_eq!(key.index(), i);
400            match i {
401                0 => assert_eq!(*value, 12),
402                1 => assert_eq!(*value, 33),
403                _ => panic!(),
404            }
405        }
406
407        i = 2;
408        for (key, value) in m.iter_mut().rev() {
409            i -= 1;
410            assert_eq!(key.index(), i);
411            match i {
412                0 => assert_eq!(*value, 12),
413                1 => assert_eq!(*value, 33),
414                _ => panic!(),
415            }
416        }
417    }
418    #[test]
419    fn keys() {
420        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
421        m.push(12);
422        m.push(33);
423
424        let mut i = 0;
425        for key in m.keys() {
426            assert_eq!(key.index(), i);
427            i += 1;
428        }
429    }
430
431    #[test]
432    fn keys_rev() {
433        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
434        m.push(12);
435        m.push(33);
436
437        let mut i = 2;
438        for key in m.keys().rev() {
439            i -= 1;
440            assert_eq!(key.index(), i);
441        }
442    }
443
444    #[test]
445    fn values() {
446        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
447        m.push(12);
448        m.push(33);
449
450        let mut i = 0;
451        for value in m.values() {
452            match i {
453                0 => assert_eq!(*value, 12),
454                1 => assert_eq!(*value, 33),
455                _ => panic!(),
456            }
457            i += 1;
458        }
459        i = 0;
460        for value_mut in m.values_mut() {
461            match i {
462                0 => assert_eq!(*value_mut, 12),
463                1 => assert_eq!(*value_mut, 33),
464                _ => panic!(),
465            }
466            i += 1;
467        }
468    }
469
470    #[test]
471    fn values_rev() {
472        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
473        m.push(12);
474        m.push(33);
475
476        let mut i = 2;
477        for value in m.values().rev() {
478            i -= 1;
479            match i {
480                0 => assert_eq!(*value, 12),
481                1 => assert_eq!(*value, 33),
482                _ => panic!(),
483            }
484        }
485        i = 2;
486        for value_mut in m.values_mut().rev() {
487            i -= 1;
488            match i {
489                0 => assert_eq!(*value_mut, 12),
490                1 => assert_eq!(*value_mut, 33),
491                _ => panic!(),
492            }
493        }
494    }
495
496    #[test]
497    fn from_iter() {
498        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
499        m.push(12);
500        m.push(33);
501
502        let n = m.values().collect::<PrimaryMap<E, _>>();
503        assert!(m.len() == n.len());
504        for (me, ne) in m.values().zip(n.values()) {
505            assert!(*me == **ne);
506        }
507    }
508
509    #[test]
510    fn from_vec() {
511        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
512        m.push(12);
513        m.push(33);
514
515        let n = PrimaryMap::<E, &usize>::from(m.values().collect::<Vec<_>>());
516        assert!(m.len() == n.len());
517        for (me, ne) in m.values().zip(n.values()) {
518            assert!(*me == **ne);
519        }
520    }
521
522    #[test]
523    fn get_many_mut() {
524        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
525        let _0 = m.push(0);
526        let _1 = m.push(1);
527        let _2 = m.push(2);
528
529        assert_eq!([&mut 0, &mut 2], m.get_disjoint_mut([_0, _2]).unwrap());
530        assert_eq!(
531            m.get_disjoint_mut([_0, _0]),
532            Err(slice::GetDisjointMutError::OverlappingIndices)
533        );
534        assert_eq!(
535            m.get_disjoint_mut([E(4)]),
536            Err(slice::GetDisjointMutError::IndexOutOfBounds)
537        );
538    }
539}
cranelift_entity/primary.rs

cranelift_entity/
primary.rs
