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    /// Analog of `get_raw` except that a raw pointer is returned rather than a
219    /// mutable reference.
220    ///
221    /// The default accessors of items in [`PrimaryMap`] will invalidate all
222    /// previous borrows obtained from the map according to miri. This function
223    /// can be used to acquire a pointer and then subsequently acquire a second
224    /// pointer later on without invalidating the first one. In other words
225    /// this is only here to help borrow two elements simultaneously with miri.
226    pub fn get_raw_mut(&mut self, k: K) -> Option<*mut V> {
227        if k.index() < self.elems.len() {
228            // SAFETY: the `add` function requires that the index is in-bounds
229            // with respect to the allocation which is satisfied here due to
230            // the bounds-check above.
231            unsafe { Some(self.elems.as_mut_ptr().add(k.index())) }
232        } else {
233            None
234        }
235    }
236}
237
238#[derive(Debug, PartialEq, Eq, Clone, Copy)]
239pub enum GetManyMutError<K> {
240    DoesNotExist(K),
241    MultipleOf(K),
242}
243
244impl<K, V> Default for PrimaryMap<K, V>
245where
246    K: EntityRef,
247{
248    fn default() -> PrimaryMap<K, V> {
249        PrimaryMap::new()
250    }
251}
252
253/// Immutable indexing into an `PrimaryMap`.
254/// The indexed value must be in the map.
255impl<K, V> Index<K> for PrimaryMap<K, V>
256where
257    K: EntityRef,
258{
259    type Output = V;
260
261    fn index(&self, k: K) -> &V {
262        &self.elems[k.index()]
263    }
264}
265
266/// Mutable indexing into an `PrimaryMap`.
267impl<K, V> IndexMut<K> for PrimaryMap<K, V>
268where
269    K: EntityRef,
270{
271    fn index_mut(&mut self, k: K) -> &mut V {
272        &mut self.elems[k.index()]
273    }
274}
275
276impl<K, V> IntoIterator for PrimaryMap<K, V>
277where
278    K: EntityRef,
279{
280    type Item = (K, V);
281    type IntoIter = IntoIter<K, V>;
282
283    fn into_iter(self) -> Self::IntoIter {
284        IntoIter::new(self.elems.into_iter())
285    }
286}
287
288impl<'a, K, V> IntoIterator for &'a PrimaryMap<K, V>
289where
290    K: EntityRef,
291{
292    type Item = (K, &'a V);
293    type IntoIter = Iter<'a, K, V>;
294
295    fn into_iter(self) -> Self::IntoIter {
296        Iter::new(self.elems.iter())
297    }
298}
299
300impl<'a, K, V> IntoIterator for &'a mut PrimaryMap<K, V>
301where
302    K: EntityRef,
303{
304    type Item = (K, &'a mut V);
305    type IntoIter = IterMut<'a, K, V>;
306
307    fn into_iter(self) -> Self::IntoIter {
308        IterMut::new(self.elems.iter_mut())
309    }
310}
311
312impl<K, V> FromIterator<V> for PrimaryMap<K, V>
313where
314    K: EntityRef,
315{
316    fn from_iter<T>(iter: T) -> Self
317    where
318        T: IntoIterator<Item = V>,
319    {
320        Self {
321            elems: Vec::from_iter(iter),
322            unused: PhantomData,
323        }
324    }
325}
326
327impl<K, V> From<Vec<V>> for PrimaryMap<K, V>
328where
329    K: EntityRef,
330{
331    fn from(elems: Vec<V>) -> Self {
332        Self {
333            elems,
334            unused: PhantomData,
335        }
336    }
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    // `EntityRef` impl for testing.
344    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
345    struct E(u32);
346
347    impl EntityRef for E {
348        fn new(i: usize) -> Self {
349            E(i as u32)
350        }
351        fn index(self) -> usize {
352            self.0 as usize
353        }
354    }
355
356    #[test]
357    fn basic() {
358        let r0 = E(0);
359        let r1 = E(1);
360        let m = PrimaryMap::<E, isize>::new();
361
362        let v: Vec<E> = m.keys().collect();
363        assert_eq!(v, []);
364
365        assert!(!m.is_valid(r0));
366        assert!(!m.is_valid(r1));
367    }
368
369    #[test]
370    fn push() {
371        let mut m = PrimaryMap::new();
372        let k0: E = m.push(12);
373        let k1 = m.push(33);
374
375        assert_eq!(m[k0], 12);
376        assert_eq!(m[k1], 33);
377
378        let v: Vec<E> = m.keys().collect();
379        assert_eq!(v, [k0, k1]);
380    }
381
382    #[test]
383    fn iter() {
384        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
385        m.push(12);
386        m.push(33);
387
388        let mut i = 0;
389        for (key, value) in &m {
390            assert_eq!(key.index(), i);
391            match i {
392                0 => assert_eq!(*value, 12),
393                1 => assert_eq!(*value, 33),
394                _ => panic!(),
395            }
396            i += 1;
397        }
398        i = 0;
399        for (key_mut, value_mut) in m.iter_mut() {
400            assert_eq!(key_mut.index(), i);
401            match i {
402                0 => assert_eq!(*value_mut, 12),
403                1 => assert_eq!(*value_mut, 33),
404                _ => panic!(),
405            }
406            i += 1;
407        }
408    }
409
410    #[test]
411    fn iter_rev() {
412        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
413        m.push(12);
414        m.push(33);
415
416        let mut i = 2;
417        for (key, value) in m.iter().rev() {
418            i -= 1;
419            assert_eq!(key.index(), i);
420            match i {
421                0 => assert_eq!(*value, 12),
422                1 => assert_eq!(*value, 33),
423                _ => panic!(),
424            }
425        }
426
427        i = 2;
428        for (key, value) in m.iter_mut().rev() {
429            i -= 1;
430            assert_eq!(key.index(), i);
431            match i {
432                0 => assert_eq!(*value, 12),
433                1 => assert_eq!(*value, 33),
434                _ => panic!(),
435            }
436        }
437    }
438    #[test]
439    fn keys() {
440        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
441        m.push(12);
442        m.push(33);
443
444        let mut i = 0;
445        for key in m.keys() {
446            assert_eq!(key.index(), i);
447            i += 1;
448        }
449    }
450
451    #[test]
452    fn keys_rev() {
453        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
454        m.push(12);
455        m.push(33);
456
457        let mut i = 2;
458        for key in m.keys().rev() {
459            i -= 1;
460            assert_eq!(key.index(), i);
461        }
462    }
463
464    #[test]
465    fn values() {
466        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
467        m.push(12);
468        m.push(33);
469
470        let mut i = 0;
471        for value in m.values() {
472            match i {
473                0 => assert_eq!(*value, 12),
474                1 => assert_eq!(*value, 33),
475                _ => panic!(),
476            }
477            i += 1;
478        }
479        i = 0;
480        for value_mut in m.values_mut() {
481            match i {
482                0 => assert_eq!(*value_mut, 12),
483                1 => assert_eq!(*value_mut, 33),
484                _ => panic!(),
485            }
486            i += 1;
487        }
488    }
489
490    #[test]
491    fn values_rev() {
492        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
493        m.push(12);
494        m.push(33);
495
496        let mut i = 2;
497        for value in m.values().rev() {
498            i -= 1;
499            match i {
500                0 => assert_eq!(*value, 12),
501                1 => assert_eq!(*value, 33),
502                _ => panic!(),
503            }
504        }
505        i = 2;
506        for value_mut in m.values_mut().rev() {
507            i -= 1;
508            match i {
509                0 => assert_eq!(*value_mut, 12),
510                1 => assert_eq!(*value_mut, 33),
511                _ => panic!(),
512            }
513        }
514    }
515
516    #[test]
517    fn from_iter() {
518        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
519        m.push(12);
520        m.push(33);
521
522        let n = m.values().collect::<PrimaryMap<E, _>>();
523        assert!(m.len() == n.len());
524        for (me, ne) in m.values().zip(n.values()) {
525            assert!(*me == **ne);
526        }
527    }
528
529    #[test]
530    fn from_vec() {
531        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
532        m.push(12);
533        m.push(33);
534
535        let n = PrimaryMap::<E, &usize>::from(m.values().collect::<Vec<_>>());
536        assert!(m.len() == n.len());
537        for (me, ne) in m.values().zip(n.values()) {
538            assert!(*me == **ne);
539        }
540    }
541
542    #[test]
543    fn get_many_mut() {
544        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
545        let _0 = m.push(0);
546        let _1 = m.push(1);
547        let _2 = m.push(2);
548
549        assert_eq!([&mut 0, &mut 2], m.get_many_mut([_0, _2]).unwrap());
550        assert_eq!(
551            m.get_many_mut([_0, _0]),
552            Err(GetManyMutError::MultipleOf(_0))
553        );
554        assert_eq!(
555            m.get_many_mut([E(4)]),
556            Err(GetManyMutError::DoesNotExist(E(4)))
557        );
558    }
559}