miri/borrow_tracker/tree_borrows/
unimap.rs

1//! This module implements the `UniMap`, which is a way to get efficient mappings
2//! optimized for the setting of `tree_borrows/tree.rs`.
3//!
4//! A `UniKeyMap<K>` is a (slow) mapping from `K` to `UniIndex`,
5//! and `UniValMap<V>` is a (fast) mapping from `UniIndex` to `V`.
6//! Thus a pair `(UniKeyMap<K>, UniValMap<V>)` acts as a virtual `HashMap<K, V>`.
7//!
8//! Because of the asymmetry in access time, the use-case for `UniMap` is the following:
9//! a tuple `(UniKeyMap<K>, Vec<UniValMap<V>>)` is much more efficient than
10//! the equivalent `Vec<HashMap<K, V>>` it represents if all maps have similar
11//! sets of keys.
12
13#![allow(dead_code)]
14
15use std::hash::Hash;
16use std::mem;
17
18use rustc_data_structures::fx::FxHashMap;
19
20/// Intermediate key between a UniKeyMap and a UniValMap.
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct UniIndex {
23    idx: u32,
24}
25
26/// From K to UniIndex
27#[derive(Debug, Clone, Default)]
28pub struct UniKeyMap<K> {
29    /// Underlying map that does all the hard work.
30    /// Key invariant: the contents of `deassigned` are disjoint from the
31    /// keys of `mapping`, and together they form the set of contiguous integers
32    /// `0 .. (mapping.len() + deassigned.len())`.
33    mapping: FxHashMap<K, u32>,
34    /// Indexes that can be reused: memory gain when the map gets sparse
35    /// due to many deletions.
36    deassigned: Vec<u32>,
37}
38
39/// From UniIndex to V
40#[derive(Debug, Clone, Eq)]
41pub struct UniValMap<V> {
42    /// The mapping data. Thanks to Vec we get both fast accesses, and
43    /// a memory-optimal representation if there are few deletions.
44    data: Vec<Option<V>>,
45}
46
47impl<V: PartialEq> UniValMap<V> {
48    /// Exact equality of two maps.
49    /// Less accurate but faster than `equivalent`, mostly because
50    /// of the fast path when the lengths are different.
51    pub fn identical(&self, other: &Self) -> bool {
52        self.data == other.data
53    }
54
55    /// Equality up to trailing `None`s of two maps, i.e.
56    /// do they represent the same mapping ?
57    pub fn equivalent(&self, other: &Self) -> bool {
58        let min_len = self.data.len().min(other.data.len());
59        self.data[min_len..].iter().all(Option::is_none)
60            && other.data[min_len..].iter().all(Option::is_none)
61            && (self.data[..min_len] == other.data[..min_len])
62    }
63}
64
65impl<V: PartialEq> PartialEq for UniValMap<V> {
66    /// 2023-05: We found that using `equivalent` rather than `identical`
67    /// in the equality testing of the `RangeMap` is neutral for most
68    /// benchmarks, while being quite beneficial for `zip-equal`
69    /// and to a lesser extent for `unicode`, `slice-get-unchecked` and
70    /// `backtraces` as well.
71    fn eq(&self, other: &Self) -> bool {
72        self.equivalent(other)
73    }
74}
75
76impl<V> Default for UniValMap<V> {
77    fn default() -> Self {
78        Self { data: Vec::default() }
79    }
80}
81
82impl<K> UniKeyMap<K>
83where
84    K: Hash + Eq,
85{
86    /// How many keys/index pairs are currently active.
87    pub fn len(&self) -> usize {
88        self.mapping.len()
89    }
90
91    /// Whether this key has an associated index or not.
92    pub fn contains_key(&self, key: &K) -> bool {
93        self.mapping.contains_key(key)
94    }
95
96    /// Assign this key to a new index. Panics if the key is already assigned,
97    /// use `get_or_insert` for a version that instead returns the existing
98    /// assignment.
99    #[track_caller]
100    pub fn insert(&mut self, key: K) -> UniIndex {
101        // We want an unused index. First we attempt to find one from `deassigned`,
102        // and if `deassigned` is empty we generate a fresh index.
103        let idx = self.deassigned.pop().unwrap_or_else(|| {
104            // `deassigned` is empty, so all keys in use are already in `mapping`.
105            // The next available key is `mapping.len()`.
106            self.mapping.len().try_into().expect("UniMap ran out of useable keys")
107        });
108        if self.mapping.insert(key, idx).is_some() {
109            panic!(
110                "This key is already assigned to a different index; either use `get_or_insert` instead if you care about this data, or first call `remove` to undo the preexisting assignment."
111            );
112        };
113        UniIndex { idx }
114    }
115
116    /// If it exists, the index this key maps to.
117    pub fn get(&self, key: &K) -> Option<UniIndex> {
118        self.mapping.get(key).map(|&idx| UniIndex { idx })
119    }
120
121    /// Either get a previously existing entry, or create a new one if it
122    /// is not yet present.
123    pub fn get_or_insert(&mut self, key: K) -> UniIndex {
124        self.get(&key).unwrap_or_else(|| self.insert(key))
125    }
126
127    /// Return whatever index this key was using to the deassigned pool.
128    ///
129    /// Note: calling this function can be dangerous. If the index still exists
130    /// somewhere in a `UniValMap` and is reassigned by the `UniKeyMap` then
131    /// it will inherit the old value of a completely unrelated key.
132    /// If you `UniKeyMap::remove` a key you should make sure to also `UniValMap::remove`
133    /// the associated `UniIndex` from ALL `UniValMap`s.
134    ///
135    /// Example of such behavior:
136    /// ```rust,ignore (private type can't be doctested)
137    /// let mut keymap = UniKeyMap::<char>::default();
138    /// let mut valmap = UniValMap::<char>::default();
139    /// // Insert 'a' -> _ -> 'A'
140    /// let idx_a = keymap.insert('a');
141    /// valmap.insert(idx_a, 'A');
142    /// // Remove 'a' -> _, but forget to remove _ -> 'A'
143    /// keymap.remove(&'a');
144    /// // valmap.remove(idx_a); // If we uncomment this line the issue is fixed
145    /// // Insert 'b' -> _
146    /// let idx_b = keymap.insert('b');
147    /// let val_b = valmap.get(idx_b);
148    /// assert_eq!(val_b, Some('A')); // Oh no
149    /// // assert_eq!(val_b, None); // This is what we would have expected
150    /// ```
151    pub fn remove(&mut self, key: &K) {
152        if let Some(idx) = self.mapping.remove(key) {
153            self.deassigned.push(idx);
154        }
155    }
156}
157
158impl<V> UniValMap<V> {
159    /// Whether this index has an associated value.
160    pub fn contains_idx(&self, idx: UniIndex) -> bool {
161        self.data.get(idx.idx as usize).and_then(Option::as_ref).is_some()
162    }
163
164    /// Reserve enough space to insert the value at the right index.
165    fn extend_to_length(&mut self, len: usize) {
166        if len > self.data.len() {
167            let nb = len - self.data.len();
168            self.data.reserve(nb);
169            for _ in 0..nb {
170                self.data.push(None);
171            }
172        }
173    }
174
175    /// Assign a value to the index. Permanently overwrites any previous value.
176    pub fn insert(&mut self, idx: UniIndex, val: V) {
177        self.extend_to_length(idx.idx as usize + 1);
178        self.data[idx.idx as usize] = Some(val)
179    }
180
181    /// Get the value at this index, if it exists.
182    pub fn get(&self, idx: UniIndex) -> Option<&V> {
183        self.data.get(idx.idx as usize).and_then(Option::as_ref)
184    }
185
186    /// Get the value at this index mutably, if it exists.
187    pub fn get_mut(&mut self, idx: UniIndex) -> Option<&mut V> {
188        self.data.get_mut(idx.idx as usize).and_then(Option::as_mut)
189    }
190
191    /// Delete any value associated with this index.
192    /// Returns None if the value was not present, otherwise
193    /// returns the previously stored value.
194    pub fn remove(&mut self, idx: UniIndex) -> Option<V> {
195        if idx.idx as usize >= self.data.len() {
196            return None;
197        }
198        let mut res = None;
199        mem::swap(&mut res, &mut self.data[idx.idx as usize]);
200        res
201    }
202}
203
204/// An access to a single value of the map.
205pub struct UniEntry<'a, V> {
206    inner: &'a mut Option<V>,
207}
208
209impl<'a, V> UniValMap<V> {
210    /// Get a wrapper around a mutable access to the value corresponding to `idx`.
211    pub fn entry(&'a mut self, idx: UniIndex) -> UniEntry<'a, V> {
212        self.extend_to_length(idx.idx as usize + 1);
213        UniEntry { inner: &mut self.data[idx.idx as usize] }
214    }
215}
216
217impl<'a, V> UniEntry<'a, V> {
218    /// Insert in the map and get the value.
219    pub fn or_insert(&mut self, default: V) -> &mut V {
220        if self.inner.is_none() {
221            *self.inner = Some(default);
222        }
223        self.inner.as_mut().unwrap()
224    }
225
226    pub fn get(&self) -> Option<&V> {
227        self.inner.as_ref()
228    }
229}
230
231mod tests {
232    use super::*;
233
234    #[test]
235    fn extend_to_length() {
236        let mut km = UniValMap::<char>::default();
237        km.extend_to_length(10);
238        assert!(km.data.len() == 10);
239        km.extend_to_length(0);
240        assert!(km.data.len() == 10);
241        km.extend_to_length(10);
242        assert!(km.data.len() == 10);
243        km.extend_to_length(11);
244        assert!(km.data.len() == 11);
245    }
246
247    #[derive(Default)]
248    struct MapWitness<K, V> {
249        key: UniKeyMap<K>,
250        val: UniValMap<V>,
251        map: FxHashMap<K, V>,
252    }
253
254    impl<K, V> MapWitness<K, V>
255    where
256        K: Copy + Hash + Eq,
257        V: Copy + Eq + std::fmt::Debug,
258    {
259        fn insert(&mut self, k: K, v: V) {
260            // UniMap
261            let i = self.key.get_or_insert(k);
262            self.val.insert(i, v);
263            // HashMap
264            self.map.insert(k, v);
265            // Consistency: nothing to check
266        }
267
268        fn get(&self, k: &K) {
269            // UniMap
270            let v1 = self.key.get(k).and_then(|i| self.val.get(i));
271            // HashMap
272            let v2 = self.map.get(k);
273            // Consistency
274            assert_eq!(v1, v2);
275        }
276
277        fn get_mut(&mut self, k: &K) {
278            // UniMap
279            let v1 = self.key.get(k).and_then(|i| self.val.get_mut(i));
280            // HashMap
281            let v2 = self.map.get_mut(k);
282            // Consistency
283            assert_eq!(v1, v2);
284        }
285        fn remove(&mut self, k: &K) {
286            // UniMap
287            if let Some(i) = self.key.get(k) {
288                self.val.remove(i);
289            }
290            self.key.remove(k);
291            // HashMap
292            self.map.remove(k);
293            // Consistency: nothing to check
294        }
295    }
296
297    #[test]
298    fn consistency_small() {
299        let mut m = MapWitness::<u64, char>::default();
300        m.insert(1, 'a');
301        m.insert(2, 'b');
302        m.get(&1);
303        m.get_mut(&2);
304        m.remove(&2);
305        m.insert(1, 'c');
306        m.get(&1);
307        m.insert(3, 'd');
308        m.insert(4, 'e');
309        m.insert(4, 'f');
310        m.get(&2);
311        m.get(&3);
312        m.get(&4);
313        m.get(&5);
314        m.remove(&100);
315        m.get_mut(&100);
316        m.get(&100);
317    }
318
319    #[test]
320    fn consistency_large() {
321        use std::collections::hash_map::DefaultHasher;
322        use std::hash::{Hash, Hasher};
323        let mut hasher = DefaultHasher::new();
324        let mut map = MapWitness::<u64, u64>::default();
325        for i in 0..1000 {
326            i.hash(&mut hasher);
327            let rng = hasher.finish();
328            let op = rng % 3 == 0;
329            let key = (rng / 2) % 50;
330            let val = (rng / 100) % 1000;
331            if op {
332                map.insert(key, val);
333            } else {
334                map.get(&key);
335            }
336        }
337    }
338}