1use crate::ctxhash::{CtxEq, CtxHash, CtxHashMap};
13use smallvec::{SmallVec, smallvec};
14
15struct Val<V> {
16 value: V,
17 level: u32,
18 generation: u32,
19}
20
21pub struct OccupiedEntry<'a, K: 'a, V: 'a> {
23 entry: crate::ctxhash::OccupiedEntry<'a, K, Val<V>>,
24}
25
26impl<'a, K, V> OccupiedEntry<'a, K, V> {
27 pub fn get(&self) -> &V {
29 &self.entry.get().value
30 }
31}
32
33pub struct VacantEntry<'a, K: 'a, V: 'a> {
35 entry: InsertLoc<'a, K, V>,
36 depth: u32,
37 generation: u32,
38}
39
40enum 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
47impl<'a, K, V> VacantEntry<'a, K, V> {
48 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
66pub enum Entry<'a, K: 'a, V: 'a> {
70 Occupied(OccupiedEntry<'a, K, V>),
71 Vacant(VacantEntry<'a, K, V>),
72}
73
74pub struct ScopedHashMap<K, V> {
80 map: CtxHashMap<K, Val<V>>,
81 generation_by_depth: SmallVec<[u32; 8]>,
82 generation: u32,
83}
84
85impl<K, V> ScopedHashMap<K, V>
86where
87 K: Clone,
88{
89 #[cfg(test)]
91 pub fn new() -> Self {
92 Self::with_capacity(16)
93 }
94
95 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 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 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 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 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 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 }
178 }
179 }
180
181 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 pub fn increment_depth(&mut self) {
198 self.generation_by_depth.push(self.generation);
199 }
200
201 pub fn decrement_depth(&mut self) {
203 self.generation += 1;
204 self.generation_by_depth.pop();
205 }
206
207 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)]
217mod tests {
218 use super::*;
219 use crate::ctxhash::NullCtx;
220
221 #[test]
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]
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}