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