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