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::map::SecondaryMap;
11use crate::EntityRef;
12use alloc::vec::Vec;
13use core::mem;
14use core::slice;
15use core::u32;
16
17#[cfg(feature = "enable-serde")]
18use serde::{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
206/// Iterating over the elements of a set.
207impl<'a, K, V> IntoIterator for &'a SparseMap<K, V>
208where
209    K: EntityRef,
210    V: SparseMapValue<K>,
211{
212    type Item = &'a V;
213    type IntoIter = slice::Iter<'a, V>;
214
215    fn into_iter(self) -> Self::IntoIter {
216        self.values()
217    }
218}
219
220/// Any `EntityRef` can be used as a sparse map value representing itself.
221impl<T> SparseMapValue<T> for T
222where
223    T: EntityRef,
224{
225    fn key(&self) -> Self {
226        *self
227    }
228}
229
230/// A sparse set of entity references.
231///
232/// Any type that implements `EntityRef` can be used as a sparse set value too.
233pub type SparseSet<T> = SparseMap<T, T>;
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238    use crate::EntityRef;
239
240    /// An opaque reference to an instruction in a function.
241    #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
242    pub struct Inst(u32);
243    entity_impl!(Inst, "inst");
244
245    // Mock key-value object for testing.
246    #[derive(PartialEq, Eq, Debug)]
247    struct Obj(Inst, &'static str);
248
249    impl SparseMapValue<Inst> for Obj {
250        fn key(&self) -> Inst {
251            self.0
252        }
253    }
254
255    #[test]
256    fn empty_immutable_map() {
257        let i1 = Inst::new(1);
258        let map: SparseMap<Inst, Obj> = SparseMap::new();
259
260        assert!(map.is_empty());
261        assert_eq!(map.len(), 0);
262        assert_eq!(map.get(i1), None);
263        assert_eq!(map.values().count(), 0);
264    }
265
266    #[test]
267    fn single_entry() {
268        let i0 = Inst::new(0);
269        let i1 = Inst::new(1);
270        let i2 = Inst::new(2);
271        let mut map = SparseMap::new();
272
273        assert!(map.is_empty());
274        assert_eq!(map.len(), 0);
275        assert_eq!(map.get(i1), None);
276        assert_eq!(map.get_mut(i1), None);
277        assert_eq!(map.remove(i1), None);
278
279        assert_eq!(map.insert(Obj(i1, "hi")), None);
280        assert!(!map.is_empty());
281        assert_eq!(map.len(), 1);
282        assert_eq!(map.get(i0), None);
283        assert_eq!(map.get(i1), Some(&Obj(i1, "hi")));
284        assert_eq!(map.get(i2), None);
285        assert_eq!(map.get_mut(i0), None);
286        assert_eq!(map.get_mut(i1), Some(&mut Obj(i1, "hi")));
287        assert_eq!(map.get_mut(i2), None);
288
289        assert_eq!(map.remove(i0), None);
290        assert_eq!(map.remove(i2), None);
291        assert_eq!(map.remove(i1), Some(Obj(i1, "hi")));
292        assert_eq!(map.len(), 0);
293        assert_eq!(map.get(i1), None);
294        assert_eq!(map.get_mut(i1), None);
295        assert_eq!(map.remove(i0), None);
296        assert_eq!(map.remove(i1), None);
297        assert_eq!(map.remove(i2), None);
298    }
299
300    #[test]
301    fn multiple_entries() {
302        let i0 = Inst::new(0);
303        let i1 = Inst::new(1);
304        let i2 = Inst::new(2);
305        let i3 = Inst::new(3);
306        let mut map = SparseMap::new();
307
308        assert_eq!(map.insert(Obj(i2, "foo")), None);
309        assert_eq!(map.insert(Obj(i1, "bar")), None);
310        assert_eq!(map.insert(Obj(i0, "baz")), None);
311
312        // Iteration order = insertion order when nothing has been removed yet.
313        assert_eq!(
314            map.values().map(|obj| obj.1).collect::<Vec<_>>(),
315            ["foo", "bar", "baz"]
316        );
317
318        assert_eq!(map.len(), 3);
319        assert_eq!(map.get(i0), Some(&Obj(i0, "baz")));
320        assert_eq!(map.get(i1), Some(&Obj(i1, "bar")));
321        assert_eq!(map.get(i2), Some(&Obj(i2, "foo")));
322        assert_eq!(map.get(i3), None);
323
324        // Remove front object, causing back to be swapped down.
325        assert_eq!(map.remove(i1), Some(Obj(i1, "bar")));
326        assert_eq!(map.len(), 2);
327        assert_eq!(map.get(i0), Some(&Obj(i0, "baz")));
328        assert_eq!(map.get(i1), None);
329        assert_eq!(map.get(i2), Some(&Obj(i2, "foo")));
330        assert_eq!(map.get(i3), None);
331
332        // Reinsert something at a previously used key.
333        assert_eq!(map.insert(Obj(i1, "barbar")), None);
334        assert_eq!(map.len(), 3);
335        assert_eq!(map.get(i0), Some(&Obj(i0, "baz")));
336        assert_eq!(map.get(i1), Some(&Obj(i1, "barbar")));
337        assert_eq!(map.get(i2), Some(&Obj(i2, "foo")));
338        assert_eq!(map.get(i3), None);
339
340        // Replace an entry.
341        assert_eq!(map.insert(Obj(i0, "bazbaz")), Some(Obj(i0, "baz")));
342        assert_eq!(map.len(), 3);
343        assert_eq!(map.get(i0), Some(&Obj(i0, "bazbaz")));
344        assert_eq!(map.get(i1), Some(&Obj(i1, "barbar")));
345        assert_eq!(map.get(i2), Some(&Obj(i2, "foo")));
346        assert_eq!(map.get(i3), None);
347
348        // Check the reference `IntoIter` impl.
349        let mut v = Vec::new();
350        for i in &map {
351            v.push(i.1);
352        }
353        assert_eq!(v.len(), map.len());
354    }
355
356    #[test]
357    fn entity_set() {
358        let i0 = Inst::new(0);
359        let i1 = Inst::new(1);
360        let mut set = SparseSet::new();
361
362        assert_eq!(set.insert(i0), None);
363        assert_eq!(set.insert(i0), Some(i0));
364        assert_eq!(set.insert(i1), None);
365        assert_eq!(set.get(i0), Some(&i0));
366        assert_eq!(set.get(i1), Some(&i1));
367    }
368}
cranelift_entity/sparse.rs

cranelift_entity/
sparse.rs
