cranelift_codegen/
scoped_hash_map.rs

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
7use crate::fx::FxHashMap;
8use core::hash::Hash;
9use smallvec::{smallvec, SmallVec};
10
11#[cfg(not(feature = "std"))]
12use crate::fx::FxHasher;
13#[cfg(not(feature = "std"))]
14type Hasher = core::hash::BuildHasherDefault<FxHasher>;
15
16struct 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.
23pub struct OccupiedEntry<'a, K: 'a, V: 'a> {
24    entry: super::hash_map::OccupiedEntry<'a, K, Val<V>>,
25}
26
27impl<'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.
35pub 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.
43enum 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
48impl<'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`.
70pub 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.
80pub struct ScopedHashMap<K, V> {
81    map: FxHashMap<K, Val<V>>,
82    generation_by_depth: SmallVec<[u32; 8]>,
83    generation: u32,
84}
85
86impl<K, V> ScopedHashMap<K, V>
87where
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)]
207mod 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}
cranelift_codegen/scoped_hash_map.rs

cranelift_codegen/
scoped_hash_map.rs
