1 //! `ScopedHashMap` 2 //! 3 //! This module defines a struct `ScopedHashMap<K, V>` which defines a `FxHashMap`-like 4 //! container that has a concept of scopes that can be entered and exited, such that 5 //! values inserted while inside a scope aren't visible outside the scope. 6 7 use core::hash::Hash; 8 use rustc_hash::FxHashMap; 9 use smallvec::{smallvec, SmallVec}; 10 11 #[cfg(not(feature = "std"))] 12 use crate::fx::FxHasher; 13 #[cfg(not(feature = "std"))] 14 type Hasher = core::hash::BuildHasherDefault<FxHasher>; 15 16 struct Val<V> { 17 value: V, 18 level: u32, 19 generation: u32, 20 } 21 22 /// A view into an occupied entry in a `ScopedHashMap`. It is part of the `Entry` enum. 23 pub struct OccupiedEntry<'a, K: 'a, V: 'a> { 24 entry: super::hash_map::OccupiedEntry<'a, K, Val<V>>, 25 } 26 27 impl<'a, K, V> OccupiedEntry<'a, K, V> { 28 /// Gets a reference to the value in the entry. 29 pub fn get(&self) -> &V { 30 &self.entry.get().value 31 } 32 } 33 34 /// A view into a vacant entry in a `ScopedHashMap`. It is part of the `Entry` enum. 35 pub struct VacantEntry<'a, K: 'a, V: 'a> { 36 entry: InsertLoc<'a, K, V>, 37 depth: u32, 38 generation: u32, 39 } 40 41 /// Where to insert from a `VacantEntry`. May be vacant or occupied in 42 /// the underlying map because of lazy (generation-based) deletion. 43 enum InsertLoc<'a, K: 'a, V: 'a> { 44 Vacant(super::hash_map::VacantEntry<'a, K, Val<V>>), 45 Occupied(super::hash_map::OccupiedEntry<'a, K, Val<V>>), 46 } 47 48 impl<'a, K, V> VacantEntry<'a, K, V> { 49 /// Sets the value of the entry with the `VacantEntry`'s key. 50 pub fn insert(self, value: V) { 51 let val = Val { 52 value, 53 level: self.depth, 54 generation: self.generation, 55 }; 56 match self.entry { 57 InsertLoc::Vacant(v) => { 58 v.insert(val); 59 } 60 InsertLoc::Occupied(mut o) => { 61 o.insert(val); 62 } 63 } 64 } 65 } 66 67 /// A view into a single entry in a map, which may either be vacant or occupied. 68 /// 69 /// This enum is constructed from the `entry` method on `ScopedHashMap`. 70 pub enum Entry<'a, K: 'a, V: 'a> { 71 Occupied(OccupiedEntry<'a, K, V>), 72 Vacant(VacantEntry<'a, K, V>), 73 } 74 75 /// A wrapper around a `FxHashMap` which adds the concept of scopes. Items inserted 76 /// within a scope are removed when the scope is exited. 77 /// 78 /// Shadowing, where one scope has entries with the same keys as a containing scope, 79 /// is not supported in this implementation. 80 pub struct ScopedHashMap<K, V> { 81 map: FxHashMap<K, Val<V>>, 82 generation_by_depth: SmallVec<[u32; 8]>, 83 generation: u32, 84 } 85 86 impl<K, V> ScopedHashMap<K, V> 87 where 88 K: PartialEq + Eq + Hash + Clone, 89 { 90 /// Creates an empty `ScopedHashMap`. 91 pub fn new() -> Self { 92 Self { 93 map: FxHashMap::default(), 94 generation: 0, 95 generation_by_depth: smallvec![0], 96 } 97 } 98 99 /// Creates an empty `ScopedHashMap` with some pre-allocated capacity. 100 pub fn with_capacity(cap: usize) -> Self { 101 let mut map = FxHashMap::default(); 102 map.reserve(cap); 103 Self { 104 map, 105 generation: 0, 106 generation_by_depth: smallvec![0], 107 } 108 } 109 110 /// Similar to `FxHashMap::entry`, gets the given key's corresponding entry in the map for 111 /// in-place manipulation. 112 pub fn entry<'a>(&'a mut self, key: K) -> Entry<'a, K, V> { 113 self.entry_with_depth(key, self.depth()) 114 } 115 116 /// Get the entry, setting the scope depth at which to insert. 117 pub fn entry_with_depth<'a>(&'a mut self, key: K, depth: usize) -> Entry<'a, K, V> { 118 debug_assert!(depth <= self.generation_by_depth.len()); 119 let generation = self.generation_by_depth[depth]; 120 let depth = depth as u32; 121 use super::hash_map::Entry::*; 122 match self.map.entry(key) { 123 Occupied(entry) => { 124 let entry_generation = entry.get().generation; 125 let entry_depth = entry.get().level as usize; 126 if self.generation_by_depth.get(entry_depth).cloned() == Some(entry_generation) { 127 Entry::Occupied(OccupiedEntry { entry }) 128 } else { 129 Entry::Vacant(VacantEntry { 130 entry: InsertLoc::Occupied(entry), 131 depth, 132 generation, 133 }) 134 } 135 } 136 Vacant(entry) => Entry::Vacant(VacantEntry { 137 entry: InsertLoc::Vacant(entry), 138 depth, 139 generation, 140 }), 141 } 142 } 143 144 /// Get a value from a key, if present. 145 pub fn get<'a>(&'a self, key: &K) -> Option<&'a V> { 146 self.map 147 .get(key) 148 .filter(|entry| { 149 let level = entry.level as usize; 150 self.generation_by_depth.get(level).cloned() == Some(entry.generation) 151 }) 152 .map(|entry| &entry.value) 153 } 154 155 /// Insert a key-value pair if absent. No-op if already exists. 156 pub fn insert_if_absent(&mut self, key: K, value: V) { 157 self.insert_if_absent_with_depth(key, value, self.depth()); 158 } 159 160 /// Insert a key-value pair if absent, using the given depth for 161 /// the insertion. No-op if already exists. 162 pub fn insert_if_absent_with_depth(&mut self, key: K, value: V, depth: usize) { 163 match self.entry_with_depth(key, depth) { 164 Entry::Vacant(v) => { 165 v.insert(value); 166 } 167 Entry::Occupied(_) => { 168 // Nothing. 169 } 170 } 171 } 172 173 /// Enter a new scope. 174 pub fn increment_depth(&mut self) { 175 self.generation_by_depth.push(self.generation); 176 } 177 178 /// Exit the current scope. 179 pub fn decrement_depth(&mut self) { 180 self.generation += 1; 181 self.generation_by_depth.pop(); 182 } 183 184 /// Return the current scope depth. 185 pub fn depth(&self) -> usize { 186 self.generation_by_depth 187 .len() 188 .checked_sub(1) 189 .expect("generation_by_depth cannot be empty") 190 } 191 } 192 193 #[cfg(test)] 194 mod tests { 195 use super::*; 196 197 #[test] 198 fn basic() { 199 let mut map: ScopedHashMap<i32, i32> = ScopedHashMap::new(); 200 201 match map.entry(0) { 202 Entry::Occupied(_entry) => panic!(), 203 Entry::Vacant(entry) => entry.insert(1), 204 } 205 match map.entry(2) { 206 Entry::Occupied(_entry) => panic!(), 207 Entry::Vacant(entry) => entry.insert(8), 208 } 209 match map.entry(2) { 210 Entry::Occupied(entry) => assert!(*entry.get() == 8), 211 Entry::Vacant(_entry) => panic!(), 212 } 213 map.increment_depth(); 214 match map.entry(2) { 215 Entry::Occupied(entry) => assert!(*entry.get() == 8), 216 Entry::Vacant(_entry) => panic!(), 217 } 218 match map.entry(1) { 219 Entry::Occupied(_entry) => panic!(), 220 Entry::Vacant(entry) => entry.insert(3), 221 } 222 match map.entry(1) { 223 Entry::Occupied(entry) => assert!(*entry.get() == 3), 224 Entry::Vacant(_entry) => panic!(), 225 } 226 match map.entry(0) { 227 Entry::Occupied(entry) => assert!(*entry.get() == 1), 228 Entry::Vacant(_entry) => panic!(), 229 } 230 match map.entry(2) { 231 Entry::Occupied(entry) => assert!(*entry.get() == 8), 232 Entry::Vacant(_entry) => panic!(), 233 } 234 map.decrement_depth(); 235 match map.entry(0) { 236 Entry::Occupied(entry) => assert!(*entry.get() == 1), 237 Entry::Vacant(_entry) => panic!(), 238 } 239 match map.entry(2) { 240 Entry::Occupied(entry) => assert!(*entry.get() == 8), 241 Entry::Vacant(_entry) => panic!(), 242 } 243 map.increment_depth(); 244 match map.entry(2) { 245 Entry::Occupied(entry) => assert!(*entry.get() == 8), 246 Entry::Vacant(_entry) => panic!(), 247 } 248 match map.entry(1) { 249 Entry::Occupied(_entry) => panic!(), 250 Entry::Vacant(entry) => entry.insert(4), 251 } 252 match map.entry(1) { 253 Entry::Occupied(entry) => assert!(*entry.get() == 4), 254 Entry::Vacant(_entry) => panic!(), 255 } 256 match map.entry(2) { 257 Entry::Occupied(entry) => assert!(*entry.get() == 8), 258 Entry::Vacant(_entry) => panic!(), 259 } 260 map.decrement_depth(); 261 map.increment_depth(); 262 map.increment_depth(); 263 map.increment_depth(); 264 match map.entry(2) { 265 Entry::Occupied(entry) => assert!(*entry.get() == 8), 266 Entry::Vacant(_entry) => panic!(), 267 } 268 match map.entry(1) { 269 Entry::Occupied(_entry) => panic!(), 270 Entry::Vacant(entry) => entry.insert(5), 271 } 272 match map.entry(1) { 273 Entry::Occupied(entry) => assert!(*entry.get() == 5), 274 Entry::Vacant(_entry) => panic!(), 275 } 276 match map.entry(2) { 277 Entry::Occupied(entry) => assert!(*entry.get() == 8), 278 Entry::Vacant(_entry) => panic!(), 279 } 280 map.decrement_depth(); 281 map.decrement_depth(); 282 map.decrement_depth(); 283 match map.entry(2) { 284 Entry::Occupied(entry) => assert!(*entry.get() == 8), 285 Entry::Vacant(_entry) => panic!(), 286 } 287 match map.entry(1) { 288 Entry::Occupied(_entry) => panic!(), 289 Entry::Vacant(entry) => entry.insert(3), 290 } 291 } 292 293 #[test] 294 fn insert_arbitrary_depth() { 295 let mut map: ScopedHashMap<i32, i32> = ScopedHashMap::new(); 296 map.insert_if_absent(1, 2); 297 assert_eq!(map.get(&1), Some(&2)); 298 map.increment_depth(); 299 assert_eq!(map.get(&1), Some(&2)); 300 map.insert_if_absent(3, 4); 301 assert_eq!(map.get(&3), Some(&4)); 302 map.decrement_depth(); 303 assert_eq!(map.get(&3), None); 304 map.increment_depth(); 305 map.insert_if_absent_with_depth(3, 4, 0); 306 assert_eq!(map.get(&3), Some(&4)); 307 map.decrement_depth(); 308 assert_eq!(map.get(&3), Some(&4)); 309 } 310 } 311