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::iter::FromIterator;
9use core::marker::PhantomData;
10use core::ops::{Index, IndexMut};
11use core::slice;
12#[cfg(feature = "enable-serde")]
13use serde::{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    /// Performs a binary search on the values with a key extraction function.
162    ///
163    /// Assumes that the values are sorted by the key extracted by the function.
164    ///
165    /// If the value is found then `Ok(K)` is returned, containing the entity key
166    /// of the matching value.
167    ///
168    /// If there are multiple matches, then any one of the matches could be returned.
169    ///
170    /// If the value is not found then Err(K) is returned, containing the entity key
171    /// where a matching element could be inserted while maintaining sorted order.
172    pub fn binary_search_values_by_key<'a, B, F>(&'a self, b: &B, f: F) -> Result<K, K>
173    where
174        F: FnMut(&'a V) -> B,
175        B: Ord,
176    {
177        self.elems
178            .binary_search_by_key(b, f)
179            .map(|i| K::new(i))
180            .map_err(|i| K::new(i))
181    }
182}
183
184impl<K, V> Default for PrimaryMap<K, V>
185where
186    K: EntityRef,
187{
188    fn default() -> PrimaryMap<K, V> {
189        PrimaryMap::new()
190    }
191}
192
193/// Immutable indexing into an `PrimaryMap`.
194/// The indexed value must be in the map.
195impl<K, V> Index<K> for PrimaryMap<K, V>
196where
197    K: EntityRef,
198{
199    type Output = V;
200
201    fn index(&self, k: K) -> &V {
202        &self.elems[k.index()]
203    }
204}
205
206/// Mutable indexing into an `PrimaryMap`.
207impl<K, V> IndexMut<K> for PrimaryMap<K, V>
208where
209    K: EntityRef,
210{
211    fn index_mut(&mut self, k: K) -> &mut V {
212        &mut self.elems[k.index()]
213    }
214}
215
216impl<K, V> IntoIterator for PrimaryMap<K, V>
217where
218    K: EntityRef,
219{
220    type Item = (K, V);
221    type IntoIter = IntoIter<K, V>;
222
223    fn into_iter(self) -> Self::IntoIter {
224        IntoIter::new(self.elems.into_iter())
225    }
226}
227
228impl<'a, K, V> IntoIterator for &'a PrimaryMap<K, V>
229where
230    K: EntityRef,
231{
232    type Item = (K, &'a V);
233    type IntoIter = Iter<'a, K, V>;
234
235    fn into_iter(self) -> Self::IntoIter {
236        Iter::new(self.elems.iter())
237    }
238}
239
240impl<'a, K, V> IntoIterator for &'a mut PrimaryMap<K, V>
241where
242    K: EntityRef,
243{
244    type Item = (K, &'a mut V);
245    type IntoIter = IterMut<'a, K, V>;
246
247    fn into_iter(self) -> Self::IntoIter {
248        IterMut::new(self.elems.iter_mut())
249    }
250}
251
252impl<K, V> FromIterator<V> for PrimaryMap<K, V>
253where
254    K: EntityRef,
255{
256    fn from_iter<T>(iter: T) -> Self
257    where
258        T: IntoIterator<Item = V>,
259    {
260        Self {
261            elems: Vec::from_iter(iter),
262            unused: PhantomData,
263        }
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270
271    // `EntityRef` impl for testing.
272    #[derive(Clone, Copy, Debug, PartialEq, Eq)]
273    struct E(u32);
274
275    impl EntityRef for E {
276        fn new(i: usize) -> Self {
277            E(i as u32)
278        }
279        fn index(self) -> usize {
280            self.0 as usize
281        }
282    }
283
284    #[test]
285    fn basic() {
286        let r0 = E(0);
287        let r1 = E(1);
288        let m = PrimaryMap::<E, isize>::new();
289
290        let v: Vec<E> = m.keys().collect();
291        assert_eq!(v, []);
292
293        assert!(!m.is_valid(r0));
294        assert!(!m.is_valid(r1));
295    }
296
297    #[test]
298    fn push() {
299        let mut m = PrimaryMap::new();
300        let k0: E = m.push(12);
301        let k1 = m.push(33);
302
303        assert_eq!(m[k0], 12);
304        assert_eq!(m[k1], 33);
305
306        let v: Vec<E> = m.keys().collect();
307        assert_eq!(v, [k0, k1]);
308    }
309
310    #[test]
311    fn iter() {
312        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
313        m.push(12);
314        m.push(33);
315
316        let mut i = 0;
317        for (key, value) in &m {
318            assert_eq!(key.index(), i);
319            match i {
320                0 => assert_eq!(*value, 12),
321                1 => assert_eq!(*value, 33),
322                _ => panic!(),
323            }
324            i += 1;
325        }
326        i = 0;
327        for (key_mut, value_mut) in m.iter_mut() {
328            assert_eq!(key_mut.index(), i);
329            match i {
330                0 => assert_eq!(*value_mut, 12),
331                1 => assert_eq!(*value_mut, 33),
332                _ => panic!(),
333            }
334            i += 1;
335        }
336    }
337
338    #[test]
339    fn iter_rev() {
340        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
341        m.push(12);
342        m.push(33);
343
344        let mut i = 2;
345        for (key, value) in m.iter().rev() {
346            i -= 1;
347            assert_eq!(key.index(), i);
348            match i {
349                0 => assert_eq!(*value, 12),
350                1 => assert_eq!(*value, 33),
351                _ => panic!(),
352            }
353        }
354
355        i = 2;
356        for (key, value) in m.iter_mut().rev() {
357            i -= 1;
358            assert_eq!(key.index(), i);
359            match i {
360                0 => assert_eq!(*value, 12),
361                1 => assert_eq!(*value, 33),
362                _ => panic!(),
363            }
364        }
365    }
366    #[test]
367    fn keys() {
368        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
369        m.push(12);
370        m.push(33);
371
372        let mut i = 0;
373        for key in m.keys() {
374            assert_eq!(key.index(), i);
375            i += 1;
376        }
377    }
378
379    #[test]
380    fn keys_rev() {
381        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
382        m.push(12);
383        m.push(33);
384
385        let mut i = 2;
386        for key in m.keys().rev() {
387            i -= 1;
388            assert_eq!(key.index(), i);
389        }
390    }
391
392    #[test]
393    fn values() {
394        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
395        m.push(12);
396        m.push(33);
397
398        let mut i = 0;
399        for value in m.values() {
400            match i {
401                0 => assert_eq!(*value, 12),
402                1 => assert_eq!(*value, 33),
403                _ => panic!(),
404            }
405            i += 1;
406        }
407        i = 0;
408        for value_mut in m.values_mut() {
409            match i {
410                0 => assert_eq!(*value_mut, 12),
411                1 => assert_eq!(*value_mut, 33),
412                _ => panic!(),
413            }
414            i += 1;
415        }
416    }
417
418    #[test]
419    fn values_rev() {
420        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
421        m.push(12);
422        m.push(33);
423
424        let mut i = 2;
425        for value in m.values().rev() {
426            i -= 1;
427            match i {
428                0 => assert_eq!(*value, 12),
429                1 => assert_eq!(*value, 33),
430                _ => panic!(),
431            }
432        }
433        i = 2;
434        for value_mut in m.values_mut().rev() {
435            i -= 1;
436            match i {
437                0 => assert_eq!(*value_mut, 12),
438                1 => assert_eq!(*value_mut, 33),
439                _ => panic!(),
440            }
441        }
442    }
443
444    #[test]
445    fn from_iter() {
446        let mut m: PrimaryMap<E, usize> = PrimaryMap::new();
447        m.push(12);
448        m.push(33);
449
450        let n = m.values().collect::<PrimaryMap<E, _>>();
451        assert!(m.len() == n.len());
452        for (me, ne) in m.values().zip(n.values()) {
453            assert!(*me == **ne);
454        }
455    }
456}