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 crate::fx::FxHashMap; 8 use core::hash::Hash; 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(), 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 /// Remote an entry. 193 pub fn remove(&mut self, key: &K) -> Option<V> { 194 self.map.remove(key).and_then(|val| { 195 let entry_generation = val.generation; 196 let entry_depth = val.level as usize; 197 if self.generation_by_depth.get(entry_depth).cloned() == Some(entry_generation) { 198 Some(val.value) 199 } else { 200 None 201 } 202 }) 203 } 204 } 205 206 #[cfg(test)] 207 mod tests { 208 use super::*; 209 210 #[test] 211 fn basic() { 212 let mut map: ScopedHashMap<i32, i32> = ScopedHashMap::new(); 213 214 match map.entry(0) { 215 Entry::Occupied(_entry) => panic!(), 216 Entry::Vacant(entry) => entry.insert(1), 217 } 218 match map.entry(2) { 219 Entry::Occupied(_entry) => panic!(), 220 Entry::Vacant(entry) => entry.insert(8), 221 } 222 match map.entry(2) { 223 Entry::Occupied(entry) => assert!(*entry.get() == 8), 224 Entry::Vacant(_entry) => panic!(), 225 } 226 map.increment_depth(); 227 match map.entry(2) { 228 Entry::Occupied(entry) => assert!(*entry.get() == 8), 229 Entry::Vacant(_entry) => panic!(), 230 } 231 match map.entry(1) { 232 Entry::Occupied(_entry) => panic!(), 233 Entry::Vacant(entry) => entry.insert(3), 234 } 235 match map.entry(1) { 236 Entry::Occupied(entry) => assert!(*entry.get() == 3), 237 Entry::Vacant(_entry) => panic!(), 238 } 239 match map.entry(0) { 240 Entry::Occupied(entry) => assert!(*entry.get() == 1), 241 Entry::Vacant(_entry) => panic!(), 242 } 243 match map.entry(2) { 244 Entry::Occupied(entry) => assert!(*entry.get() == 8), 245 Entry::Vacant(_entry) => panic!(), 246 } 247 map.decrement_depth(); 248 match map.entry(0) { 249 Entry::Occupied(entry) => assert!(*entry.get() == 1), 250 Entry::Vacant(_entry) => panic!(), 251 } 252 match map.entry(2) { 253 Entry::Occupied(entry) => assert!(*entry.get() == 8), 254 Entry::Vacant(_entry) => panic!(), 255 } 256 map.increment_depth(); 257 match map.entry(2) { 258 Entry::Occupied(entry) => assert!(*entry.get() == 8), 259 Entry::Vacant(_entry) => panic!(), 260 } 261 match map.entry(1) { 262 Entry::Occupied(_entry) => panic!(), 263 Entry::Vacant(entry) => entry.insert(4), 264 } 265 match map.entry(1) { 266 Entry::Occupied(entry) => assert!(*entry.get() == 4), 267 Entry::Vacant(_entry) => panic!(), 268 } 269 match map.entry(2) { 270 Entry::Occupied(entry) => assert!(*entry.get() == 8), 271 Entry::Vacant(_entry) => panic!(), 272 } 273 map.decrement_depth(); 274 map.increment_depth(); 275 map.increment_depth(); 276 map.increment_depth(); 277 match map.entry(2) { 278 Entry::Occupied(entry) => assert!(*entry.get() == 8), 279 Entry::Vacant(_entry) => panic!(), 280 } 281 match map.entry(1) { 282 Entry::Occupied(_entry) => panic!(), 283 Entry::Vacant(entry) => entry.insert(5), 284 } 285 match map.entry(1) { 286 Entry::Occupied(entry) => assert!(*entry.get() == 5), 287 Entry::Vacant(_entry) => panic!(), 288 } 289 match map.entry(2) { 290 Entry::Occupied(entry) => assert!(*entry.get() == 8), 291 Entry::Vacant(_entry) => panic!(), 292 } 293 map.decrement_depth(); 294 map.decrement_depth(); 295 map.decrement_depth(); 296 match map.entry(2) { 297 Entry::Occupied(entry) => assert!(*entry.get() == 8), 298 Entry::Vacant(_entry) => panic!(), 299 } 300 match map.entry(1) { 301 Entry::Occupied(_entry) => panic!(), 302 Entry::Vacant(entry) => entry.insert(3), 303 } 304 } 305 306 #[test] 307 fn insert_arbitrary_depth() { 308 let mut map: ScopedHashMap<i32, i32> = ScopedHashMap::new(); 309 map.insert_if_absent(1, 2); 310 assert_eq!(map.get(&1), Some(&2)); 311 map.increment_depth(); 312 assert_eq!(map.get(&1), Some(&2)); 313 map.insert_if_absent(3, 4); 314 assert_eq!(map.get(&3), Some(&4)); 315 map.decrement_depth(); 316 assert_eq!(map.get(&3), None); 317 map.increment_depth(); 318 map.insert_if_absent_with_depth(3, 4, 0); 319 assert_eq!(map.get(&3), Some(&4)); 320 map.decrement_depth(); 321 assert_eq!(map.get(&3), Some(&4)); 322 } 323 } 324