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}