cranelift_entity/
sparse.rs

1//! Sparse mapping of entity references to larger value types.
2//!
3//! This module provides a `SparseMap` data structure which implements a sparse mapping from an
4//! `EntityRef` key to a value type that may be on the larger side. This implementation is based on
5//! the paper:
6//!
7//! > Briggs, Torczon, *An efficient representation for sparse sets*,
8//! > ACM Letters on Programming Languages and Systems, Volume 2, Issue 1-4, March-Dec. 1993.
9
10use crate::EntityRef;
11use crate::map::SecondaryMap;
12use alloc::vec::Vec;
13use core::fmt;
14use core::mem;
15use core::slice;
16
17#[cfg(feature = "enable-serde")]
18use serde_derive::{Deserialize, Serialize};
19
20/// Trait for extracting keys from values stored in a `SparseMap`.
21///
22/// All values stored in a `SparseMap` must keep track of their own key in the map and implement
23/// this trait to provide access to the key.
24pub trait SparseMapValue<K> {
25    /// Get the key of this sparse map value. This key is not allowed to change while the value
26    /// is a member of the map.
27    fn key(&self) -> K;
28}
29
30/// A sparse mapping of entity references.
31///
32/// A `SparseMap<K, V>` map provides:
33///
34/// - Memory usage equivalent to `SecondaryMap<K, u32>` + `Vec<V>`, so much smaller than
35///   `SecondaryMap<K, V>` for sparse mappings of larger `V` types.
36/// - Constant time lookup, slightly slower than `SecondaryMap`.
37/// - A very fast, constant time `clear()` operation.
38/// - Fast insert and erase operations.
39/// - Stable iteration that is as fast as a `Vec<V>`.
40///
41/// # Compared to `SecondaryMap`
42///
43/// When should we use a `SparseMap` instead of a secondary `SecondaryMap`? First of all,
44/// `SparseMap` does not provide the functionality of a `PrimaryMap` which can allocate and assign
45/// entity references to objects as they are pushed onto the map. It is only the secondary entity
46/// maps that can be replaced with a `SparseMap`.
47///
48/// - A secondary entity map assigns a default mapping to all keys. It doesn't distinguish between
49///   an unmapped key and one that maps to the default value. `SparseMap` does not require
50///   `Default` values, and it tracks accurately if a key has been mapped or not.
51/// - Iterating over the contents of an `SecondaryMap` is linear in the size of the *key space*,
52///   while iterating over a `SparseMap` is linear in the number of elements in the mapping. This
53///   is an advantage precisely when the mapping is sparse.
54/// - `SparseMap::clear()` is constant time and super-fast. `SecondaryMap::clear()` is linear in
55///   the size of the key space. (Or, rather the required `resize()` call following the `clear()`
56///   is).
57/// - `SparseMap` requires the values to implement `SparseMapValue<K>` which means that they must
58///   contain their own key.
59#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
60pub struct SparseMap<K, V>
61where
62    K: EntityRef,
63    V: SparseMapValue<K>,
64{
65    sparse: SecondaryMap<K, u32>,
66    dense: Vec<V>,
67}
68
69impl<K, V> SparseMap<K, V>
70where
71    K: EntityRef,
72    V: SparseMapValue<K>,
73{
74    /// Create a new empty mapping.
75    pub fn new() -> Self {
76        Self {
77            sparse: SecondaryMap::new(),
78            dense: Vec::new(),
79        }
80    }
81
82    /// Returns the number of elements in the map.
83    pub fn len(&self) -> usize {
84        self.dense.len()
85    }
86
87    /// Returns true is the map contains no elements.
88    pub fn is_empty(&self) -> bool {
89        self.dense.is_empty()
90    }
91
92    /// Remove all elements from the mapping.
93    pub fn clear(&mut self) {
94        self.dense.clear();
95    }
96
97    /// Returns a reference to the value corresponding to the key.
98    pub fn get(&self, key: K) -> Option<&V> {
99        if let Some(idx) = self.sparse.get(key).cloned() {
100            if let Some(entry) = self.dense.get(idx as usize) {
101                if entry.key() == key {
102                    return Some(entry);
103                }
104            }
105        }
106        None
107    }
108
109    /// Returns a mutable reference to the value corresponding to the key.
110    ///
111    /// Note that the returned value must not be mutated in a way that would change its key. This
112    /// would invalidate the sparse set data structure.
113    pub fn get_mut(&mut self, key: K) -> Option<&mut V> {
114        if let Some(idx) = self.sparse.get(key).cloned() {
115            if let Some(entry) = self.dense.get_mut(idx as usize) {
116                if entry.key() == key {
117                    return Some(entry);
118                }
119            }
120        }
121        None
122    }
123
124    /// Return the index into `dense` of the value corresponding to `key`.
125    fn index(&self, key: K) -> Option<usize> {
126        if let Some(idx) = self.sparse.get(key).cloned() {
127            let idx = idx as usize;
128            if let Some(entry) = self.dense.get(idx) {
129                if entry.key() == key {
130                    return Some(idx);
131                }
132            }
133        }
134        None
135    }
136
137    /// Return `true` if the map contains a value corresponding to `key`.
138    pub fn contains_key(&self, key: K) -> bool {
139        self.get(key).is_some()
140    }
141
142    /// Insert a value into the map.
143    ///
144    /// If the map did not have this key present, `None` is returned.
145    ///
146    /// If the map did have this key present, the value is updated, and the old value is returned.
147    ///
148    /// It is not necessary to provide a key since the value knows its own key already.
149    pub fn insert(&mut self, value: V) -> Option<V> {
150        let key = value.key();
151
152        // Replace the existing entry for `key` if there is one.
153        if let Some(entry) = self.get_mut(key) {
154            return Some(mem::replace(entry, value));
155        }
156
157        // There was no previous entry for `key`. Add it to the end of `dense`.
158        let idx = self.dense.len();
159        debug_assert!(idx <= u32::MAX as usize, "SparseMap overflow");
160        self.dense.push(value);
161        self.sparse[key] = idx as u32;
162        None
163    }
164
165    /// Remove a value from the map and return it.
166    pub fn remove(&mut self, key: K) -> Option<V> {
167        if let Some(idx) = self.index(key) {
168            let back = self.dense.pop().unwrap();
169
170            // Are we popping the back of `dense`?
171            if idx == self.dense.len() {
172                return Some(back);
173            }
174
175            // We're removing an element from the middle of `dense`.
176            // Replace the element at `idx` with the back of `dense`.
177            // Repair `sparse` first.
178            self.sparse[back.key()] = idx as u32;
179            return Some(mem::replace(&mut self.dense[idx], back));
180        }
181
182        // Nothing to remove.
183        None
184    }
185
186    /// Remove the last value from the map.
187    pub fn pop(&mut self) -> Option<V> {
188        self.dense.pop()
189    }
190
191    /// Get an iterator over the values in the map.
192    ///
193    /// The iteration order is entirely determined by the preceding sequence of `insert` and
194    /// `remove` operations. In particular, if no elements were removed, this is the insertion
195    /// order.
196    pub fn values(&self) -> slice::Iter<'_, V> {
197        self.dense.iter()
198    }
199
200    /// Get the values as a slice.
201    pub fn as_slice(&self) -> &[V] {
202        self.dense.as_slice()
203    }
204}
205
206impl<K, V> Default for SparseMap<K, V>
207where
208    K: EntityRef,
209    V: SparseMapValue<K>,
210{
211    fn default() -> SparseMap<K, V> {
212        SparseMap::new()
213    }
214}
215
216impl<K, V> fmt::Debug for SparseMap<K, V>
217where
218    K: EntityRef + fmt::Debug,
219    V: SparseMapValue<K> + fmt::Debug,
220{
221    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
222        f.debug_map()
223            .entries(self.values().map(|v| (v.key(), v)))
224            .finish()
225    }
226}
227
228/// Iterating over the elements of a set.
229impl<'a, K, V> IntoIterator for &'a SparseMap<K, V>
230where
231    K: EntityRef,
232    V: SparseMapValue<K>,
233{
234    type Item = &'a V;
235    type IntoIter = slice::Iter<'a, V>;
236
237    fn into_iter(self) -> Self::IntoIter {
238        self.values()
239    }
240}
241
242/// Any `EntityRef` can be used as a sparse map value representing itself.
243impl<T> SparseMapValue<T> for T
244where
245    T: EntityRef,
246{
247    fn key(&self) -> Self {
248        *self
249    }
250}
251
252/// A sparse set of entity references.
253///
254/// Any type that implements `EntityRef` can be used as a sparse set value too.
255pub type SparseSet<T> = SparseMap<T, T>;
256
257#[cfg(test)]
258mod tests {
259    use alloc::format;
260
261    use super::*;
262
263    /// An opaque reference to an instruction in a function.
264    #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
265    pub struct Inst(u32);
266    entity_impl!(Inst, "inst");
267
268    // Mock key-value object for testing.
269    #[derive(PartialEq, Eq, Debug)]
270    struct Obj(Inst, &'static str);
271
272    impl SparseMapValue<Inst> for Obj {
273        fn key(&self) -> Inst {
274            self.0
275        }
276    }
277
278    #[test]
279    fn empty_immutable_map() {
280        let i1 = Inst::new(1);
281        let map: SparseMap<Inst, Obj> = SparseMap::new();
282
283        assert!(map.is_empty());
284        assert_eq!(map.len(), 0);
285        assert_eq!(map.get(i1), None);
286        assert_eq!(map.values().count(), 0);
287    }
288
289    #[test]
290    fn single_entry() {
291        let i0 = Inst::new(0);
292        let i1 = Inst::new(1);
293        let i2 = Inst::new(2);
294        let mut map = SparseMap::new();
295
296        assert!(map.is_empty());
297        assert_eq!(map.len(), 0);
298        assert_eq!(map.get(i1), None);
299        assert_eq!(map.get_mut(i1), None);
300        assert_eq!(map.remove(i1), None);
301
302        assert_eq!(map.insert(Obj(i1, "hi")), None);
303        assert!(!map.is_empty());
304        assert_eq!(map.len(), 1);
305        assert_eq!(map.get(i0), None);
306        assert_eq!(map.get(i1), Some(&Obj(i1, "hi")));
307        assert_eq!(map.get(i2), None);
308        assert_eq!(map.get_mut(i0), None);
309        assert_eq!(map.get_mut(i1), Some(&mut Obj(i1, "hi")));
310        assert_eq!(map.get_mut(i2), None);
311
312        assert_eq!(map.remove(i0), None);
313        assert_eq!(map.remove(i2), None);
314        assert_eq!(map.remove(i1), Some(Obj(i1, "hi")));
315        assert_eq!(map.len(), 0);
316        assert_eq!(map.get(i1), None);
317        assert_eq!(map.get_mut(i1), None);
318        assert_eq!(map.remove(i0), None);
319        assert_eq!(map.remove(i1), None);
320        assert_eq!(map.remove(i2), None);
321    }
322
323    #[test]
324    fn multiple_entries() {
325        let i0 = Inst::new(0);
326        let i1 = Inst::new(1);
327        let i2 = Inst::new(2);
328        let i3 = Inst::new(3);
329        let mut map = SparseMap::new();
330
331        assert_eq!(map.insert(Obj(i2, "foo")), None);
332        assert_eq!(map.insert(Obj(i1, "bar")), None);
333        assert_eq!(map.insert(Obj(i0, "baz")), None);
334
335        // Iteration order = insertion order when nothing has been removed yet.
336        assert_eq!(
337            map.values().map(|obj| obj.1).collect::<Vec<_>>(),
338            ["foo", "bar", "baz"]
339        );
340
341        assert_eq!(map.len(), 3);
342        assert_eq!(map.get(i0), Some(&Obj(i0, "baz")));
343        assert_eq!(map.get(i1), Some(&Obj(i1, "bar")));
344        assert_eq!(map.get(i2), Some(&Obj(i2, "foo")));
345        assert_eq!(map.get(i3), None);
346
347        // Remove front object, causing back to be swapped down.
348        assert_eq!(map.remove(i1), Some(Obj(i1, "bar")));
349        assert_eq!(map.len(), 2);
350        assert_eq!(map.get(i0), Some(&Obj(i0, "baz")));
351        assert_eq!(map.get(i1), None);
352        assert_eq!(map.get(i2), Some(&Obj(i2, "foo")));
353        assert_eq!(map.get(i3), None);
354
355        // Reinsert something at a previously used key.
356        assert_eq!(map.insert(Obj(i1, "barbar")), None);
357        assert_eq!(map.len(), 3);
358        assert_eq!(map.get(i0), Some(&Obj(i0, "baz")));
359        assert_eq!(map.get(i1), Some(&Obj(i1, "barbar")));
360        assert_eq!(map.get(i2), Some(&Obj(i2, "foo")));
361        assert_eq!(map.get(i3), None);
362
363        // Replace an entry.
364        assert_eq!(map.insert(Obj(i0, "bazbaz")), Some(Obj(i0, "baz")));
365        assert_eq!(map.len(), 3);
366        assert_eq!(map.get(i0), Some(&Obj(i0, "bazbaz")));
367        assert_eq!(map.get(i1), Some(&Obj(i1, "barbar")));
368        assert_eq!(map.get(i2), Some(&Obj(i2, "foo")));
369        assert_eq!(map.get(i3), None);
370
371        // Check the reference `IntoIter` impl.
372        let mut v = Vec::new();
373        for i in &map {
374            v.push(i.1);
375        }
376        assert_eq!(v.len(), map.len());
377    }
378
379    #[test]
380    fn entity_set() {
381        let i0 = Inst::new(0);
382        let i1 = Inst::new(1);
383        let mut set = SparseSet::new();
384
385        assert_eq!(set.insert(i0), None);
386        assert_eq!(set.insert(i0), Some(i0));
387        assert_eq!(set.insert(i1), None);
388        assert_eq!(set.get(i0), Some(&i0));
389        assert_eq!(set.get(i1), Some(&i1));
390    }
391
392    #[test]
393    fn default_impl() {
394        let map: SparseMap<Inst, Obj> = SparseMap::default();
395
396        assert!(map.is_empty());
397        assert_eq!(map.len(), 0);
398    }
399
400    #[test]
401    fn debug_impl() {
402        let i1 = Inst::new(1);
403        let mut map = SparseMap::new();
404        assert_eq!(map.insert(Obj(i1, "hi")), None);
405
406        let debug = format!("{map:?}");
407        let expected = "{inst1: Obj(inst1, \"hi\")}";
408        assert_eq!(debug, expected);
409    }
410}
cranelift_entity/sparse.rs

cranelift_entity/
sparse.rs
