src/mono_hash_map.rs

   1 //! This is a "monotonic HashMap": A HashMap 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 HashMap, because they will not be dropped
   4 //! (from being removed) or moved (because they are boxed).
   5 //! The API is 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
   8 use std::collections::hash_map::Entry;
   9 use std::cell::RefCell;
  10 use std::hash::Hash;
  11 use std::borrow::Borrow;
  12
  13 use rustc_data_structures::fx::FxHashMap;
  14
  15 use crate::AllocMap;
  16
  17 #[derive(Debug, Clone)]
  18 pub struct MonoHashMap<K: Hash + Eq, V>(RefCell<FxHashMap<K, Box<V>>>);
  19
  20 impl<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 memory of constants does not show up in this list.
  24     pub fn iter<T>(&self, f: impl FnOnce(&mut dyn Iterator<Item=(&K, &V)>) -> T) -> T {
  25         f(&mut self.0.borrow().iter().map(|(k, v)| (k, &**v)))
  26     }
  27 }
  28
  29 impl<K: Hash + Eq, V> Default for MonoHashMap<K, V> {
  30     fn default() -> Self {
  31         MonoHashMap(RefCell::new(Default::default()))
  32     }
  33 }
  34
  35 impl<K: Hash + Eq, V> AllocMap<K, V> for MonoHashMap<K, V> {
  36     #[inline(always)]
  37     fn contains_key<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> bool
  38         where K: Borrow<Q>
  39     {
  40         self.0.get_mut().contains_key(k)
  41     }
  42
  43     #[inline(always)]
  44     fn insert(&mut self, k: K, v: V) -> Option<V>
  45     {
  46         self.0.get_mut().insert(k, Box::new(v)).map(|x| *x)
  47     }
  48
  49     #[inline(always)]
  50     fn remove<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> Option<V>
  51         where K: Borrow<Q>
  52     {
  53         self.0.get_mut().remove(k).map(|x| *x)
  54     }
  55
  56     #[inline(always)]
  57     fn filter_map_collect<T>(&self, mut f: impl FnMut(&K, &V) -> Option<T>) -> Vec<T> {
  58         self.0.borrow()
  59             .iter()
  60             .filter_map(move |(k, v)| f(k, &*v))
  61             .collect()
  62     }
  63
  64     /// The most interesting method: Providing a shared ref without
  65     /// holding the `RefCell` open, and inserting new data if the key
  66     /// is not used yet.
  67     /// `vacant` is called if the key is not found in the map;
  68     /// if it returns a reference, that is used directly, if it
  69     /// returns owned data, that is put into the map and returned.
  70     #[inline(always)]
  71     fn get_or<E>(
  72         &self,
  73         k: K,
  74         vacant: impl FnOnce() -> Result<V, E>
  75     ) -> Result<&V, E> {
  76         let val: *const V = match self.0.borrow_mut().entry(k) {
  77             Entry::Occupied(entry) => &**entry.get(),
  78             Entry::Vacant(entry) => &**entry.insert(Box::new(vacant()?)),
  79         };
  80         // This is safe because `val` points into a `Box`, that we know will not move and
  81         // will also not be dropped as long as the shared reference `self` is live.
  82         unsafe { Ok(&*val) }
  83     }
  84
  85     #[inline(always)]
  86     fn get_mut_or<E>(
  87         &mut self,
  88         k: K,
  89         vacant: impl FnOnce() -> Result<V, E>
  90     ) -> Result<&mut V, E>
  91     {
  92         match self.0.get_mut().entry(k) {
  93             Entry::Occupied(e) => Ok(e.into_mut()),
  94             Entry::Vacant(e) => {
  95                 let v = vacant()?;
  96                 Ok(e.insert(Box::new(v)))
  97             }
  98         }
  99     }
 100 }