miri/mono_hash_map.rs
1//! This is a "monotonic `FxHashMap`": A `FxHashMap` that, when shared, can be pushed to but not
2//! otherwise mutated. We also box items in the map. This means we can safely provide
3//! shared references into existing items in the `FxHashMap`, because they will not be dropped
4//! (from being removed) or moved (because they are boxed).
5//! The API is completely tailored to what `memory.rs` needs. It is still in
6//! a separate file to minimize the amount of code that has to care about the unsafety.
7
8use std::borrow::Borrow;
9use std::cell::RefCell;
10use std::collections::hash_map::Entry;
11use std::hash::Hash;
12
13use rustc_data_structures::fx::FxHashMap;
14
15use crate::AllocMap;
16
17#[derive(Debug, Clone)]
18pub struct MonoHashMap<K: Hash + Eq, V>(RefCell<FxHashMap<K, Box<V>>>);
19
20impl<K: Hash + Eq, V> MonoHashMap<K, V> {
21 /// This function exists for priroda to be able to iterate over all evaluator memory.
22 ///
23 /// The function is somewhat roundabout with the closure argument because internally the
24 /// `MonoHashMap` uses a `RefCell`. When iterating over the `FxHashMap` inside the `RefCell`,
25 /// we need to keep a borrow to the `FxHashMap` inside the iterator. The borrow is only alive
26 /// as long as the `Ref` returned by `RefCell::borrow()` is alive. So we can't return the
27 /// iterator, as that would drop the `Ref`. We can't return both, as it's not possible in Rust
28 /// to have a struct/tuple with a field that refers to another field.
29 pub fn iter<T>(&self, f: impl FnOnce(&mut dyn Iterator<Item = (&K, &V)>) -> T) -> T {
30 f(&mut self.0.borrow().iter().map(|(k, v)| (k, &**v)))
31 }
32}
33
34impl<K: Hash + Eq, V> Default for MonoHashMap<K, V> {
35 fn default() -> Self {
36 MonoHashMap(RefCell::new(Default::default()))
37 }
38}
39
40impl<K: Hash + Eq, V> AllocMap<K, V> for MonoHashMap<K, V> {
41 #[inline(always)]
42 fn contains_key<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> bool
43 where
44 K: Borrow<Q>,
45 {
46 self.0.get_mut().contains_key(k)
47 }
48
49 #[inline(always)]
50 fn contains_key_ref<Q: ?Sized + Hash + Eq>(&self, k: &Q) -> bool
51 where
52 K: Borrow<Q>,
53 {
54 self.0.borrow().contains_key(k)
55 }
56
57 #[inline(always)]
58 fn insert(&mut self, k: K, v: V) -> Option<V> {
59 self.0.get_mut().insert(k, Box::new(v)).map(|x| *x)
60 }
61
62 #[inline(always)]
63 fn remove<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> Option<V>
64 where
65 K: Borrow<Q>,
66 {
67 self.0.get_mut().remove(k).map(|x| *x)
68 }
69
70 #[inline(always)]
71 fn filter_map_collect<T>(&self, mut f: impl FnMut(&K, &V) -> Option<T>) -> Vec<T> {
72 self.0.borrow().iter().filter_map(move |(k, v)| f(k, v)).collect()
73 }
74
75 /// The most interesting method: Providing a shared reference without
76 /// holding the `RefCell` open, and inserting new data if the key
77 /// is not used yet.
78 /// `vacant` is called if the key is not found in the map;
79 /// if it returns a reference, that is used directly, if it
80 /// returns owned data, that is put into the map and returned.
81 #[inline(always)]
82 fn get_or<E>(&self, k: K, vacant: impl FnOnce() -> Result<V, E>) -> Result<&V, E> {
83 // We cannot hold borrow_mut while calling `vacant`, since that might have to do lookups in this very map.
84 if let Some(v) = self.0.borrow().get(&k) {
85 let val: *const V = &**v;
86 // This is safe because `val` points into a `Box`, that we know will not move and
87 // will also not be dropped as long as the shared reference `self` is live.
88 return unsafe { Ok(&*val) };
89 }
90 let new_val = Box::new(vacant()?);
91 let val: *const V = &**self.0.borrow_mut().try_insert(k, new_val).ok().unwrap();
92 // This is safe because `val` points into a `Box`, that we know will not move and
93 // will also not be dropped as long as the shared reference `self` is live.
94 unsafe { Ok(&*val) }
95 }
96
97 /// Read-only lookup (avoid read-acquiring the RefCell).
98 fn get(&self, k: K) -> Option<&V> {
99 let val: *const V = match self.0.borrow().get(&k) {
100 Some(v) => &**v,
101 None => return None,
102 };
103 // This is safe because `val` points into a `Box`, that we know will not move and
104 // will also not be dropped as long as the shared reference `self` is live.
105 unsafe { Some(&*val) }
106 }
107
108 #[inline(always)]
109 fn get_mut_or<E>(&mut self, k: K, vacant: impl FnOnce() -> Result<V, E>) -> Result<&mut V, E> {
110 match self.0.get_mut().entry(k) {
111 Entry::Occupied(e) => Ok(e.into_mut()),
112 Entry::Vacant(e) => {
113 let v = vacant()?;
114 Ok(e.insert(Box::new(v)))
115 }
116 }
117 }
118}