diff --git a/src/tools/miri/src/borrow_tracker/tree_borrows/unimap.rs b/src/tools/miri/src/borrow_tracker/tree_borrows/unimap.rs new file mode 100644 index 000000000000..c1d452ca89e9 --- /dev/null +++ b/src/tools/miri/src/borrow_tracker/tree_borrows/unimap.rs @@ -0,0 +1,304 @@ +//! This module implements the `UniMap`, which is a way to get efficient mappings +//! optimized for the setting of `tree_borrows/tree.rs`. +//! +//! A `UniKeyMap` is a (slow) mapping from `K` to `UniIndex`, +//! and `UniValMap` is a (fast) mapping from `UniIndex` to `V`. +//! Thus a pair `(UniKeyMap, UniValMap)` acts as a virtual `HashMap`. +//! +//! Because of the asymmetry in access time, the use-case for `UniMap` is the following: +//! a tuple `(UniKeyMap, Vec>)` is much more efficient than +//! the equivalent `Vec>` it represents if all maps have similar +//! sets of keys. + +#![allow(dead_code)] + +use std::hash::Hash; + +use rustc_data_structures::fx::FxHashMap; + +/// Intermediate key between a UniKeyMap and a UniValMap. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct UniIndex { + idx: u32, +} + +/// From K to UniIndex +#[derive(Debug, Clone, Default)] +pub struct UniKeyMap { + /// Underlying map that does all the hard work. + /// Key invariant: the contents of `deassigned` are disjoint from the + /// keys of `mapping`, and together they form the set of contiguous integers + /// `0 .. (mapping.len() + deassigned.len())`. + mapping: FxHashMap, + /// Indexes that can be reused: memory gain when the map gets sparse + /// due to many deletions. + deassigned: Vec, +} + +/// From UniIndex to V +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct UniValMap { + /// The mapping data. Thanks to Vec we get both fast accesses, and + /// a memory-optimal representation if there are few deletions. + data: Vec>, +} + +impl Default for UniValMap { + fn default() -> Self { + Self { data: Vec::default() } + } +} + +impl UniKeyMap +where + K: Hash + Eq, +{ + /// How many keys/index pairs are currently active. + pub fn len(&self) -> usize { + self.mapping.len() + } + + /// Whether this key has an associated index or not. + pub fn contains_key(&self, key: &K) -> bool { + self.mapping.contains_key(key) + } + + /// Assign this key to a new index. Panics if the key is already assigned, + /// use `get_or_insert` for a version that instead returns the existing + /// assignment. + #[track_caller] + pub fn insert(&mut self, key: K) -> UniIndex { + // We want an unused index. First we attempt to find one from `deassigned`, + // and if `deassigned` is empty we generate a fresh index. + let idx = self.deassigned.pop().unwrap_or_else(|| { + // `deassigned` is empty, so all keys in use are already in `mapping`. + // The next available key is `mapping.len()`. + self.mapping.len().try_into().expect("UniMap ran out of useable keys") + }); + if self.mapping.insert(key, idx).is_some() { + panic!( + "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." + ); + }; + UniIndex { idx } + } + + /// If it exists, the index this key maps to. + pub fn get(&self, key: &K) -> Option { + self.mapping.get(key).map(|&idx| UniIndex { idx }) + } + + /// Either get a previously existing entry, or create a new one if it + /// is not yet present. + pub fn get_or_insert(&mut self, key: K) -> UniIndex { + self.get(&key).unwrap_or_else(|| self.insert(key)) + } + + /// Return whatever index this key was using to the deassigned pool. + /// + /// Note: calling this function can be dangerous. If the index still exists + /// somewhere in a `UniValMap` and is reassigned by the `UniKeyMap` then + /// it will inherit the old value of a completely unrelated key. + /// If you `UniKeyMap::remove` a key you should make sure to also `UniValMap::remove` + /// the associated `UniIndex` from ALL `UniValMap`s. + /// + /// Example of such behavior: + /// ``` + /// let mut keymap = UniKeyMap::::default(); + /// let mut valmap = UniValMap::::default(); + /// // Insert 'a' -> _ -> 'A' + /// let idx_a = keymap.insert('a'); + /// valmap.insert(idx_a, 'A'); + /// // Remove 'a' -> _, but forget to remove _ -> 'A' + /// keymap.remove(&'a'); + /// // valmap.remove(idx_a); // If we uncomment this line the issue is fixed + /// // Insert 'b' -> _ + /// let idx_b = keymap.insert('b'); + /// let val_b = valmap.get(idx_b); + /// assert_eq!(val_b, Some('A')); // Oh no + /// // assert_eq!(val_b, None); // This is what we would have expected + /// ``` + pub fn remove(&mut self, key: &K) { + if let Some(idx) = self.mapping.remove(key) { + self.deassigned.push(idx); + } + } +} + +impl UniValMap { + /// Whether this index has an associated value. + pub fn contains_idx(&self, idx: UniIndex) -> bool { + self.data.get(idx.idx as usize).and_then(Option::as_ref).is_some() + } + + /// Reserve enough space to insert the value at the right index. + fn extend_to_length(&mut self, len: usize) { + if len > self.data.len() { + let nb = len - self.data.len(); + self.data.reserve(nb); + for _ in 0..nb { + self.data.push(None); + } + } + } + + /// Assign a value to the index. Permanently overwrites any previous value. + pub fn insert(&mut self, idx: UniIndex, val: V) { + self.extend_to_length(idx.idx as usize + 1); + self.data[idx.idx as usize] = Some(val) + } + + /// Get the value at this index, if it exists. + pub fn get(&self, idx: UniIndex) -> Option<&V> { + self.data.get(idx.idx as usize).and_then(Option::as_ref) + } + + /// Get the value at this index mutably, if it exists. + pub fn get_mut(&mut self, idx: UniIndex) -> Option<&mut V> { + self.data.get_mut(idx.idx as usize).and_then(Option::as_mut) + } + + /// Delete any value associated with this index. Ok even if the index + /// has no associated value. + pub fn remove(&mut self, idx: UniIndex) { + if idx.idx as usize >= self.data.len() { + return; + } + self.data[idx.idx as usize] = None; + } +} + +/// An access to a single value of the map. +pub struct UniEntry<'a, V> { + inner: &'a mut Option, +} + +impl<'a, V> UniValMap { + /// Get a wrapper around a mutable access to the value corresponding to `idx`. + pub fn entry(&'a mut self, idx: UniIndex) -> UniEntry<'a, V> { + self.extend_to_length(idx.idx as usize + 1); + UniEntry { inner: &mut self.data[idx.idx as usize] } + } +} + +impl<'a, V> UniEntry<'a, V> { + /// Insert in the map and get the value. + pub fn or_insert_with(&mut self, default: F) -> &mut V + where + F: FnOnce() -> V, + { + if self.inner.is_none() { + *self.inner = Some(default()); + } + self.inner.as_mut().unwrap() + } +} + +mod tests { + use super::*; + + #[test] + fn extend_to_length() { + let mut km = UniValMap::::default(); + km.extend_to_length(10); + assert!(km.data.len() == 10); + km.extend_to_length(0); + assert!(km.data.len() == 10); + km.extend_to_length(10); + assert!(km.data.len() == 10); + km.extend_to_length(11); + assert!(km.data.len() == 11); + } + + #[derive(Default)] + struct MapWitness { + key: UniKeyMap, + val: UniValMap, + map: FxHashMap, + } + + impl MapWitness + where + K: Copy + Hash + Eq, + V: Copy + Eq + std::fmt::Debug, + { + fn insert(&mut self, k: K, v: V) { + // UniMap + let i = self.key.get_or_insert(k); + self.val.insert(i, v); + // HashMap + self.map.insert(k, v); + // Consistency: nothing to check + } + + fn get(&self, k: &K) { + // UniMap + let v1 = self.key.get(k).and_then(|i| self.val.get(i)); + // HashMap + let v2 = self.map.get(k); + // Consistency + assert_eq!(v1, v2); + } + + fn get_mut(&mut self, k: &K) { + // UniMap + let v1 = self.key.get(k).and_then(|i| self.val.get_mut(i)); + // HashMap + let v2 = self.map.get_mut(k); + // Consistency + assert_eq!(v1, v2); + } + fn remove(&mut self, k: &K) { + // UniMap + if let Some(i) = self.key.get(k) { + self.val.remove(i); + } + self.key.remove(k); + // HashMap + self.map.remove(k); + // Consistency: nothing to check + } + } + + #[test] + fn consistency_small() { + let mut m = MapWitness::::default(); + m.insert(1, 'a'); + m.insert(2, 'b'); + m.get(&1); + m.get_mut(&2); + m.remove(&2); + m.insert(1, 'c'); + m.get(&1); + m.insert(3, 'd'); + m.insert(4, 'e'); + m.insert(4, 'f'); + m.get(&2); + m.get(&3); + m.get(&4); + m.get(&5); + m.remove(&100); + m.get_mut(&100); + m.get(&100); + } + + #[test] + fn consistency_large() { + use std::collections::hash_map::DefaultHasher; + use std::hash::{Hash, Hasher}; + let mut hasher = DefaultHasher::new(); + let mut map = MapWitness::::default(); + for i in 0..1000 { + i.hash(&mut hasher); + let rng = hasher.finish(); + let op = rng % 3 == 0; + let key = (rng / 2) % 50; + let val = (rng / 100) % 1000; + if op { + map.insert(key, val); + } else { + map.get(&key); + } + } + } +}