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