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

miri/borrow_tracker/tree_borrows/
unimap.rs
