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> Default for MonoHashMap<K, V> {
  21     fn default() -> Self {
  22         MonoHashMap(RefCell::new(Default::default()))
  23     }
  24 }
  25
  26 impl<K: Hash + Eq, V> AllocMap<K, V> for MonoHashMap<K, V> {
  27     #[inline(always)]
  28     fn contains_key<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> bool
  29         where K: Borrow<Q>
  30     {
  31         self.0.get_mut().contains_key(k)
  32     }
  33
  34     #[inline(always)]
  35     fn insert(&mut self, k: K, v: V) -> Option<V>
  36     {
  37         self.0.get_mut().insert(k, Box::new(v)).map(|x| *x)
  38     }
  39
  40     #[inline(always)]
  41     fn remove<Q: ?Sized + Hash + Eq>(&mut self, k: &Q) -> Option<V>
  42         where K: Borrow<Q>
  43     {
  44         self.0.get_mut().remove(k).map(|x| *x)
  45     }
  46
  47     #[inline(always)]
  48     fn filter_map_collect<T>(&self, mut f: impl FnMut(&K, &V) -> Option<T>) -> Vec<T> {
  49         self.0.borrow()
  50             .iter()
  51             .filter_map(move |(k, v)| f(k, &*v))
  52             .collect()
  53     }
  54
  55     /// The most interesting method: Providing a shared ref without
  56     /// holding the `RefCell` open, and inserting new data if the key
  57     /// is not used yet.
  58     /// `vacant` is called if the key is not found in the map;
  59     /// if it returns a reference, that is used directly, if it
  60     /// returns owned data, that is put into the map and returned.
  61     #[inline(always)]
  62     fn get_or<E>(
  63         &self,
  64         k: K,
  65         vacant: impl FnOnce() -> Result<V, E>
  66     ) -> Result<&V, E> {
  67         let val: *const V = match self.0.borrow_mut().entry(k) {
  68             Entry::Occupied(entry) => &**entry.get(),
  69             Entry::Vacant(entry) => &**entry.insert(Box::new(vacant()?)),
  70         };
  71         // This is safe because `val` points into a `Box`, that we know will not move and
  72         // will also not be dropped as long as the shared reference `self` is live.
  73         unsafe { Ok(&*val) }
  74     }
  75
  76     #[inline(always)]
  77     fn get_mut_or<E>(
  78         &mut self,
  79         k: K,
  80         vacant: impl FnOnce() -> Result<V, E>
  81     ) -> Result<&mut V, E>
  82     {
  83         match self.0.get_mut().entry(k) {
  84             Entry::Occupied(e) => Ok(e.into_mut()),
  85             Entry::Vacant(e) => {
  86                 let v = vacant()?;
  87                 Ok(e.insert(Box::new(v)))
  88             }
  89         }
  90     }
  91 }