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 core::mem;
10 
11 struct Val<K, V> {
12     value: V,
13     next_key: Option<K>,
14     depth: usize,
15 }
16 
17 /// A view into an occupied entry in a `ScopedHashMap`. It is part of the `Entry` enum.
18 pub struct OccupiedEntry<'a, K: 'a, V: 'a> {
19     entry: super::hash_map::OccupiedEntry<'a, K, Val<K, V>>,
20 }
21 
22 impl<'a, K, V> OccupiedEntry<'a, K, V> {
23     /// Gets a reference to the value in the entry.
get(&self) -> &V24     pub fn get(&self) -> &V {
25         &self.entry.get().value
26     }
27 }
28 
29 /// A view into a vacant entry in a `ScopedHashMap`. It is part of the `Entry` enum.
30 pub struct VacantEntry<'a, K: 'a, V: 'a> {
31     entry: super::hash_map::VacantEntry<'a, K, Val<K, V>>,
32     next_key: Option<K>,
33     depth: usize,
34 }
35 
36 impl<'a, K, V> VacantEntry<'a, K, V> {
37     /// Sets the value of the entry with the `VacantEntry`'s key.
insert(self, value: V)38     pub fn insert(self, value: V) {
39         self.entry.insert(Val {
40             value,
41             next_key: self.next_key,
42             depth: self.depth,
43         });
44     }
45 }
46 
47 /// A view into a single entry in a map, which may either be vacant or occupied.
48 ///
49 /// This enum is constructed from the `entry` method on `ScopedHashMap`.
50 pub enum Entry<'a, K: 'a, V: 'a> {
51     Occupied(OccupiedEntry<'a, K, V>),
52     Vacant(VacantEntry<'a, K, V>),
53 }
54 
55 /// A wrapper around a `FxHashMap` which adds the concept of scopes. Items inserted
56 /// within a scope are removed when the scope is exited.
57 ///
58 /// Shadowing, where one scope has entries with the same keys as a containing scope,
59 /// is not supported in this implementation.
60 pub struct ScopedHashMap<K, V> {
61     map: FxHashMap<K, Val<K, V>>,
62     last_insert: Option<K>,
63     current_depth: usize,
64 }
65 
66 impl<K, V> ScopedHashMap<K, V>
67 where
68     K: PartialEq + Eq + Hash + Clone,
69 {
70     /// Creates an empty `ScopedHashMap`.
new() -> Self71     pub fn new() -> Self {
72         Self {
73             map: FxHashMap(),
74             last_insert: None,
75             current_depth: 0,
76         }
77     }
78 
79     /// Similar to `FxHashMap::entry`, gets the given key's corresponding entry in the map for
80     /// in-place manipulation.
entry(&mut self, key: K) -> Entry<K, V>81     pub fn entry(&mut self, key: K) -> Entry<K, V> {
82         use super::hash_map::Entry::*;
83         match self.map.entry(key) {
84             Occupied(entry) => Entry::Occupied(OccupiedEntry { entry }),
85             Vacant(entry) => {
86                 let clone_key = entry.key().clone();
87                 Entry::Vacant(VacantEntry {
88                     entry,
89                     next_key: mem::replace(&mut self.last_insert, Some(clone_key)),
90                     depth: self.current_depth,
91                 })
92             }
93         }
94     }
95 
96     /// Enter a new scope.
increment_depth(&mut self)97     pub fn increment_depth(&mut self) {
98         // Increment the depth.
99         self.current_depth = self.current_depth.checked_add(1).unwrap();
100     }
101 
102     /// Exit the current scope.
decrement_depth(&mut self)103     pub fn decrement_depth(&mut self) {
104         // Remove all elements inserted at the current depth.
105         while let Some(key) = self.last_insert.clone() {
106             use crate::hash_map::Entry::*;
107             match self.map.entry(key) {
108                 Occupied(entry) => {
109                     if entry.get().depth != self.current_depth {
110                         break;
111                     }
112                     self.last_insert = entry.remove_entry().1.next_key;
113                 }
114                 Vacant(_) => panic!(),
115             }
116         }
117 
118         // Decrement the depth.
119         self.current_depth = self.current_depth.checked_sub(1).unwrap();
120     }
121 }
122 
123 #[cfg(test)]
124 mod tests {
125     use super::*;
126 
127     #[test]
basic()128     fn basic() {
129         let mut map: ScopedHashMap<i32, i32> = ScopedHashMap::new();
130 
131         match map.entry(0) {
132             Entry::Occupied(_entry) => panic!(),
133             Entry::Vacant(entry) => entry.insert(1),
134         }
135         match map.entry(2) {
136             Entry::Occupied(_entry) => panic!(),
137             Entry::Vacant(entry) => entry.insert(8),
138         }
139         match map.entry(2) {
140             Entry::Occupied(entry) => assert!(*entry.get() == 8),
141             Entry::Vacant(_entry) => panic!(),
142         }
143         map.increment_depth();
144         match map.entry(2) {
145             Entry::Occupied(entry) => assert!(*entry.get() == 8),
146             Entry::Vacant(_entry) => panic!(),
147         }
148         match map.entry(1) {
149             Entry::Occupied(_entry) => panic!(),
150             Entry::Vacant(entry) => entry.insert(3),
151         }
152         match map.entry(1) {
153             Entry::Occupied(entry) => assert!(*entry.get() == 3),
154             Entry::Vacant(_entry) => panic!(),
155         }
156         match map.entry(0) {
157             Entry::Occupied(entry) => assert!(*entry.get() == 1),
158             Entry::Vacant(_entry) => panic!(),
159         }
160         match map.entry(2) {
161             Entry::Occupied(entry) => assert!(*entry.get() == 8),
162             Entry::Vacant(_entry) => panic!(),
163         }
164         map.decrement_depth();
165         match map.entry(0) {
166             Entry::Occupied(entry) => assert!(*entry.get() == 1),
167             Entry::Vacant(_entry) => panic!(),
168         }
169         match map.entry(2) {
170             Entry::Occupied(entry) => assert!(*entry.get() == 8),
171             Entry::Vacant(_entry) => panic!(),
172         }
173         map.increment_depth();
174         match map.entry(2) {
175             Entry::Occupied(entry) => assert!(*entry.get() == 8),
176             Entry::Vacant(_entry) => panic!(),
177         }
178         match map.entry(1) {
179             Entry::Occupied(_entry) => panic!(),
180             Entry::Vacant(entry) => entry.insert(4),
181         }
182         match map.entry(1) {
183             Entry::Occupied(entry) => assert!(*entry.get() == 4),
184             Entry::Vacant(_entry) => panic!(),
185         }
186         match map.entry(2) {
187             Entry::Occupied(entry) => assert!(*entry.get() == 8),
188             Entry::Vacant(_entry) => panic!(),
189         }
190         map.decrement_depth();
191         map.increment_depth();
192         map.increment_depth();
193         map.increment_depth();
194         match map.entry(2) {
195             Entry::Occupied(entry) => assert!(*entry.get() == 8),
196             Entry::Vacant(_entry) => panic!(),
197         }
198         match map.entry(1) {
199             Entry::Occupied(_entry) => panic!(),
200             Entry::Vacant(entry) => entry.insert(5),
201         }
202         match map.entry(1) {
203             Entry::Occupied(entry) => assert!(*entry.get() == 5),
204             Entry::Vacant(_entry) => panic!(),
205         }
206         match map.entry(2) {
207             Entry::Occupied(entry) => assert!(*entry.get() == 8),
208             Entry::Vacant(_entry) => panic!(),
209         }
210         map.decrement_depth();
211         map.decrement_depth();
212         map.decrement_depth();
213         match map.entry(2) {
214             Entry::Occupied(entry) => assert!(*entry.get() == 8),
215             Entry::Vacant(_entry) => panic!(),
216         }
217         match map.entry(1) {
218             Entry::Occupied(_entry) => panic!(),
219             Entry::Vacant(entry) => entry.insert(3),
220         }
221     }
222 }
223