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