crates/hir_def/src/intern.rs

   1 //! Global `Arc`-based object interning infrastructure.
   2 //!
   3 //! Eventually this should probably be replaced with salsa-based interning.
   4
   5 use std::{
   6     collections::HashMap,
   7     fmt::{self, Debug, Display},
   8     hash::{BuildHasherDefault, Hash, Hasher},
   9     ops::Deref,
  10     sync::Arc,
  11 };
  12
  13 use dashmap::{DashMap, SharedValue};
  14 use lock_api::RwLockWriteGuard;
  15 use once_cell::sync::OnceCell;
  16 use parking_lot::RawRwLock;
  17 use rustc_hash::FxHasher;
  18
  19 use crate::generics::GenericParams;
  20
  21 type InternMap<T> = DashMap<Arc<T>, (), BuildHasherDefault<FxHasher>>;
  22 type Guard<T> = RwLockWriteGuard<
  23     'static,
  24     RawRwLock,
  25     HashMap<Arc<T>, SharedValue<()>, BuildHasherDefault<FxHasher>>,
  26 >;
  27
  28 pub struct Interned<T: Internable + ?Sized> {
  29     arc: Arc<T>,
  30 }
  31
  32 impl<T: Internable> Interned<T> {
  33     pub fn new(obj: T) -> Self {
  34         match Interned::lookup(&obj) {
  35             Ok(this) => this,
  36             Err(shard) => {
  37                 let arc = Arc::new(obj);
  38                 Self::alloc(arc, shard)
  39             }
  40         }
  41     }
  42 }
  43
  44 impl<T: Internable + ?Sized> Interned<T> {
  45     fn lookup(obj: &T) -> Result<Self, Guard<T>> {
  46         let storage = T::storage().get();
  47         let shard_idx = storage.determine_map(obj);
  48         let shard = &storage.shards()[shard_idx];
  49         let shard = shard.write();
  50
  51         // Atomically,
  52         // - check if `obj` is already in the map
  53         //   - if so, clone its `Arc` and return it
  54         //   - if not, box it up, insert it, and return a clone
  55         // This needs to be atomic (locking the shard) to avoid races with other thread, which could
  56         // insert the same object between us looking it up and inserting it.
  57
  58         // FIXME: avoid double lookup/hashing by using raw entry API (once stable, or when
  59         // hashbrown can be plugged into dashmap)
  60         match shard.get_key_value(obj) {
  61             Some((arc, _)) => Ok(Self { arc: arc.clone() }),
  62             None => Err(shard),
  63         }
  64     }
  65
  66     fn alloc(arc: Arc<T>, mut shard: Guard<T>) -> Self {
  67         let arc2 = arc.clone();
  68
  69         shard.insert(arc2, SharedValue::new(()));
  70
  71         Self { arc }
  72     }
  73 }
  74
  75 impl Interned<str> {
  76     pub fn new_str(s: &str) -> Self {
  77         match Interned::lookup(s) {
  78             Ok(this) => this,
  79             Err(shard) => {
  80                 let arc = Arc::<str>::from(s);
  81                 Self::alloc(arc, shard)
  82             }
  83         }
  84     }
  85 }
  86
  87 impl<T: Internable + ?Sized> Drop for Interned<T> {
  88     #[inline]
  89     fn drop(&mut self) {
  90         // When the last `Ref` is dropped, remove the object from the global map.
  91         if Arc::strong_count(&self.arc) == 2 {
  92             // Only `self` and the global map point to the object.
  93
  94             self.drop_slow();
  95         }
  96     }
  97 }
  98
  99 impl<T: Internable + ?Sized> Interned<T> {
 100     #[cold]
 101     fn drop_slow(&mut self) {
 102         let storage = T::storage().get();
 103         let shard_idx = storage.determine_map(&self.arc);
 104         let shard = &storage.shards()[shard_idx];
 105         let mut shard = shard.write();
 106
 107         // FIXME: avoid double lookup
 108         let (arc, _) = shard.get_key_value(&self.arc).expect("interned value removed prematurely");
 109
 110         if Arc::strong_count(arc) != 2 {
 111             // Another thread has interned another copy
 112             return;
 113         }
 114
 115         shard.remove(&self.arc);
 116
 117         // Shrink the backing storage if the shard is less than 50% occupied.
 118         if shard.len() * 2 < shard.capacity() {
 119             shard.shrink_to_fit();
 120         }
 121     }
 122 }
 123
 124 /// Compares interned `Ref`s using pointer equality.
 125 impl<T: Internable> PartialEq for Interned<T> {
 126     // NOTE: No `?Sized` because `ptr_eq` doesn't work right with trait objects.
 127
 128     #[inline]
 129     fn eq(&self, other: &Self) -> bool {
 130         Arc::ptr_eq(&self.arc, &other.arc)
 131     }
 132 }
 133
 134 impl<T: Internable> Eq for Interned<T> {}
 135
 136 impl PartialEq for Interned<str> {
 137     fn eq(&self, other: &Self) -> bool {
 138         Arc::ptr_eq(&self.arc, &other.arc)
 139     }
 140 }
 141
 142 impl Eq for Interned<str> {}
 143
 144 impl<T: Internable + ?Sized> Hash for Interned<T> {
 145     fn hash<H: Hasher>(&self, state: &mut H) {
 146         // NOTE: Cast disposes vtable pointer / slice/str length.
 147         state.write_usize(Arc::as_ptr(&self.arc) as *const () as usize)
 148     }
 149 }
 150
 151 impl<T: Internable + ?Sized> AsRef<T> for Interned<T> {
 152     #[inline]
 153     fn as_ref(&self) -> &T {
 154         &self.arc
 155     }
 156 }
 157
 158 impl<T: Internable + ?Sized> Deref for Interned<T> {
 159     type Target = T;
 160
 161     #[inline]
 162     fn deref(&self) -> &Self::Target {
 163         &self.arc
 164     }
 165 }
 166
 167 impl<T: Internable + ?Sized> Clone for Interned<T> {
 168     fn clone(&self) -> Self {
 169         Self { arc: self.arc.clone() }
 170     }
 171 }
 172
 173 impl<T: Debug + Internable + ?Sized> Debug for Interned<T> {
 174     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 175         (*self.arc).fmt(f)
 176     }
 177 }
 178
 179 impl<T: Display + Internable + ?Sized> Display for Interned<T> {
 180     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 181         (*self.arc).fmt(f)
 182     }
 183 }
 184
 185 pub struct InternStorage<T: ?Sized> {
 186     map: OnceCell<InternMap<T>>,
 187 }
 188
 189 impl<T: ?Sized> InternStorage<T> {
 190     pub const fn new() -> Self {
 191         Self { map: OnceCell::new() }
 192     }
 193 }
 194
 195 impl<T: Internable + ?Sized> InternStorage<T> {
 196     fn get(&self) -> &InternMap<T> {
 197         self.map.get_or_init(DashMap::default)
 198     }
 199 }
 200
 201 pub trait Internable: Hash + Eq + 'static {
 202     fn storage() -> &'static InternStorage<Self>;
 203 }
 204
 205 /// Implements `Internable` for a given list of types, making them usable with `Interned`.
 206 #[macro_export]
 207 #[doc(hidden)]
 208 macro_rules! _impl_internable {
 209     ( $($t:path),+ $(,)? ) => { $(
 210         impl Internable for $t {
 211             fn storage() -> &'static InternStorage<Self> {
 212                 static STORAGE: InternStorage<$t> = InternStorage::new();
 213                 &STORAGE
 214             }
 215         }
 216     )+ };
 217 }
 218
 219 pub use crate::_impl_internable as impl_internable;
 220
 221 impl_internable!(
 222     crate::type_ref::TypeRef,
 223     crate::type_ref::TraitRef,
 224     crate::type_ref::TypeBound,
 225     crate::path::ModPath,
 226     crate::path::GenericArgs,
 227     crate::attr::AttrInput,
 228     GenericParams,
 229     str,
 230 );