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

cranelift_codegen/
scoped_hash_map.rs
