cranelift_codegen/
scoped_hash_map.rs

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
12use crate::ctxhash::{CtxEq, CtxHash, CtxHashMap};
13use smallvec::{SmallVec, smallvec};
14
15struct 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.
22pub 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    /// Gets a reference to the value in the entry.
28    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.
34pub 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.
42enum 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    /// Sets the value of the entry with the `VacantEntry`'s key.
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`.
69pub 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.
79pub 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    /// Creates an empty `ScopedHashMap`.
90    #[cfg(test)]
91    pub fn new() -> Self {
92        Self::with_capacity(16)
93    }
94
95    /// Creates an empty `ScopedHashMap` with some pre-allocated capacity.
96    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.
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.
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.
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.
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.
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.
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.
197    pub fn increment_depth(&mut self) {
198        self.generation_by_depth.push(self.generation);
199    }
200
201    /// Exit the current scope.
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.
208    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}