Global TypeRef/TraitRef interning

Author: jonas-schievink

Cargo.lock

@@ -91,7 +91,7 @@ impl HirDisplay for Function {
         let ret_type = if !qual.is_async {
             &data.ret_type
         } else {
-            match &data.ret_type {
+            match &*data.ret_type {
                 TypeRef::ImplTrait(bounds) => match &bounds[0] {
                     TypeBound::Path(path) => {
                         path.segments().iter().last().unwrap().args_and_bindings.unwrap().bindings

crates/hir/src/display.rs

@@ -957,7 +957,7 @@ impl SelfParam {
         func_data
             .params
             .first()
-            .map(|param| match *param {
+            .map(|param| match &**param {
                 TypeRef::Reference(.., mutability) => match mutability {
                     hir_def::type_ref::Mutability::Shared => Access::Shared,
                     hir_def::type_ref::Mutability::Mut => Access::Exclusive,
@@ -1011,7 +1011,7 @@ impl Const {
     }
 
     pub fn type_ref(self, db: &dyn HirDatabase) -> TypeRef {
-        db.const_data(self.id).type_ref.clone()
+        db.const_data(self.id).type_ref.as_ref().clone()
     }
 }
 
@@ -1101,7 +1101,7 @@ impl TypeAlias {
     }
 
     pub fn type_ref(self, db: &dyn HirDatabase) -> Option<TypeRef> {
-        db.type_alias_data(self.id).type_ref.clone()
+        db.type_alias_data(self.id).type_ref.as_deref().cloned()
     }
 
     pub fn ty(self, db: &dyn HirDatabase) -> Type {
@@ -1615,7 +1615,7 @@ impl Impl {
     // FIXME: the return type is wrong. This should be a hir version of
     // `TraitRef` (ie, resolved `TypeRef`).
     pub fn trait_(self, db: &dyn HirDatabase) -> Option<TraitRef> {
-        db.impl_data(self.id).target_trait.clone()
+        db.impl_data(self.id).target_trait.as_deref().cloned()
     }
 
     pub fn self_ty(self, db: &dyn HirDatabase) -> Type {

crates/hir/src/lib.rs

@@ -11,6 +11,7 @@ doctest = false
 
 [dependencies]
 cov-mark = { version = "1.1", features = ["thread-local"] }
+dashmap = { version = "4.0.2", features = ["raw-api"] }
 log = "0.4.8"
 once_cell = "1.3.1"
 rustc-hash = "1.1.0"

crates/hir_def/Cargo.toml

@@ -15,6 +15,7 @@ use tt::{Delimiter, DelimiterKind, Leaf, Subtree, TokenTree};
 use crate::{
     body::{CfgExpander, LowerCtx},
     db::DefDatabase,
+    intern::Interned,
     item_tree::{AttrOwner, Field, Fields, ItemTree, ModItem, RawVisibilityId},
     src::HasChildSource,
     src::HasSource,
@@ -58,7 +59,7 @@ pub enum VariantData {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct FieldData {
     pub name: Name,
-    pub type_ref: TypeRef,
+    pub type_ref: Interned<TypeRef>,
     pub visibility: RawVisibility,
 }
 
@@ -292,7 +293,7 @@ fn lower_struct(
                     || Either::Left(fd.clone()),
                     || FieldData {
                         name: Name::new_tuple_field(i),
-                        type_ref: TypeRef::from_ast_opt(&ctx, fd.ty()),
+                        type_ref: Interned::new(TypeRef::from_ast_opt(&ctx, fd.ty())),
                         visibility: RawVisibility::from_ast(db, ast.with_value(fd.visibility())),
                     },
                 );
@@ -309,7 +310,7 @@ fn lower_struct(
                     || Either::Right(fd.clone()),
                     || FieldData {
                         name: fd.name().map(|n| n.as_name()).unwrap_or_else(Name::missing),
-                        type_ref: TypeRef::from_ast_opt(&ctx, fd.ty()),
+                        type_ref: Interned::new(TypeRef::from_ast_opt(&ctx, fd.ty())),
                         visibility: RawVisibility::from_ast(db, ast.with_value(fd.visibility())),
                     },
                 );
@@ -358,7 +359,7 @@ fn lower_field(
 ) -> FieldData {
     FieldData {
         name: field.name.clone(),
-        type_ref: item_tree[field.type_ref].clone(),
+        type_ref: field.type_ref.clone(),
         visibility: item_tree[override_visibility.unwrap_or(field.visibility)].clone(),
     }
 }

crates/hir_def/src/adt.rs

@@ -9,6 +9,7 @@ use crate::{
     attr::Attrs,
     body::Expander,
     db::DefDatabase,
+    intern::Interned,
     item_tree::{AssocItem, FunctionQualifier, ItemTreeId, ModItem, Param},
     type_ref::{TraitRef, TypeBound, TypeRef},
     visibility::RawVisibility,
@@ -19,8 +20,8 @@ use crate::{
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct FunctionData {
     pub name: Name,
-    pub params: Vec<TypeRef>,
-    pub ret_type: TypeRef,
+    pub params: Vec<Interned<TypeRef>>,
+    pub ret_type: Interned<TypeRef>,
     pub attrs: Attrs,
     /// True if the first param is `self`. This is relevant to decide whether this
     /// can be called as a method.
@@ -57,11 +58,11 @@ impl FunctionData {
             params: enabled_params
                 .clone()
                 .filter_map(|id| match &item_tree[id] {
-                    Param::Normal(ty) => Some(item_tree[*ty].clone()),
+                    Param::Normal(ty) => Some(ty.clone()),
                     Param::Varargs => None,
                 })
                 .collect(),
-            ret_type: item_tree[func.ret_type].clone(),
+            ret_type: func.ret_type.clone(),
             attrs: item_tree.attrs(db, krate, ModItem::from(loc.id.value).into()),
             has_self_param: func.has_self_param,
             has_body: func.has_body,
@@ -76,7 +77,7 @@ impl FunctionData {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct TypeAliasData {
     pub name: Name,
-    pub type_ref: Option<TypeRef>,
+    pub type_ref: Option<Interned<TypeRef>>,
     pub visibility: RawVisibility,
     pub is_extern: bool,
     /// Bounds restricting the type alias itself (eg. `type Ty: Bound;` in a trait or impl).
@@ -94,7 +95,7 @@ impl TypeAliasData {
 
         Arc::new(TypeAliasData {
             name: typ.name.clone(),
-            type_ref: typ.type_ref.map(|id| item_tree[id].clone()),
+            type_ref: typ.type_ref.clone(),
             visibility: item_tree[typ.visibility].clone(),
             is_extern: typ.is_extern,
             bounds: typ.bounds.to_vec(),
@@ -156,8 +157,8 @@ impl TraitData {
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct ImplData {
-    pub target_trait: Option<TraitRef>,
-    pub self_ty: TypeRef,
+    pub target_trait: Option<Interned<TraitRef>>,
+    pub self_ty: Interned<TypeRef>,
     pub items: Vec<AssocItemId>,
     pub is_negative: bool,
 }
@@ -169,8 +170,8 @@ impl ImplData {
 
         let item_tree = impl_loc.id.item_tree(db);
         let impl_def = &item_tree[impl_loc.id.value];
-        let target_trait = impl_def.target_trait.map(|id| item_tree[id].clone());
-        let self_ty = item_tree[impl_def.self_ty].clone();
+        let target_trait = impl_def.target_trait.clone();
+        let self_ty = impl_def.self_ty.clone();
         let is_negative = impl_def.is_negative;
         let module_id = impl_loc.container;
         let container = AssocContainerId::ImplId(id);
@@ -195,7 +196,7 @@ impl ImplData {
 pub struct ConstData {
     /// const _: () = ();
     pub name: Option<Name>,
-    pub type_ref: TypeRef,
+    pub type_ref: Interned<TypeRef>,
     pub visibility: RawVisibility,
 }
 
@@ -207,7 +208,7 @@ impl ConstData {
 
         Arc::new(ConstData {
             name: konst.name.clone(),
-            type_ref: item_tree[konst.type_ref].clone(),
+            type_ref: konst.type_ref.clone(),
             visibility: item_tree[konst.visibility].clone(),
         })
     }
@@ -216,7 +217,7 @@ impl ConstData {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct StaticData {
     pub name: Option<Name>,
-    pub type_ref: TypeRef,
+    pub type_ref: Interned<TypeRef>,
     pub visibility: RawVisibility,
     pub mutable: bool,
     pub is_extern: bool,
@@ -230,7 +231,7 @@ impl StaticData {
 
         Arc::new(StaticData {
             name: Some(statik.name.clone()),
-            type_ref: item_tree[statik.type_ref].clone(),
+            type_ref: statik.type_ref.clone(),
             visibility: item_tree[statik.visibility].clone(),
             mutable: statik.mutable,
             is_extern: statik.is_extern,

crates/hir_def/src/data.rs

@@ -0,0 +1,157 @@
+//! Global `Arc`-based object interning infrastructure.
+//!
+//! Eventually this should probably be replaced with salsa-based interning.
+
+use std::{
+    fmt::{self, Debug},
+    hash::{BuildHasherDefault, Hash},
+    ops::Deref,
+    sync::Arc,
+};
+
+use dashmap::{DashMap, SharedValue};
+use once_cell::sync::OnceCell;
+use rustc_hash::FxHasher;
+
+type InternMap<T> = DashMap<Arc<T>, (), BuildHasherDefault<FxHasher>>;
+
+pub struct Interned<T: Internable> {
+    arc: Arc<T>,
+}
+
+impl<T: Internable> Interned<T> {
+    pub fn new(obj: T) -> Self {
+        let storage = T::storage().get();
+        let shard_idx = storage.determine_map(&obj);
+        let shard = &storage.shards()[shard_idx];
+        let shard = shard.upgradeable_read();
+
+        // Atomically,
+        // - check if `obj` is already in the map
+        //   - if so, clone its `Arc` and return it
+        //   - if not, box it up, insert it, and return a clone
+        // This needs to be atomic (locking the shard) to avoid races with other thread, which could
+        // insert the same object between us looking it up and inserting it.
+
+        // FIXME: avoid double lookup by using raw entry API (once stable, or when hashbrown can be
+        // plugged into dashmap)
+        if let Some((arc, _)) = shard.get_key_value(&obj) {
+            return Self { arc: arc.clone() };
+        }
+
+        let arc = Arc::new(obj);
+        let arc2 = arc.clone();
+
+        {
+            let mut shard = shard.upgrade();
+            shard.insert(arc2, SharedValue::new(()));
+        }
+
+        Self { arc }
+    }
+}
+
+impl<T: Internable> Drop for Interned<T> {
+    fn drop(&mut self) {
+        // When the last `Ref` is dropped, remove the object from the global map.
+        if Arc::strong_count(&self.arc) == 2 {
+            // Only `self` and the global map point to the object.
+
+            let storage = T::storage().get();
+            let shard_idx = storage.determine_map(&self.arc);
+            let shard = &storage.shards()[shard_idx];
+            let mut shard = shard.write();
+
+            // FIXME: avoid double lookup
+            let (arc, _) =
+                shard.get_key_value(&self.arc).expect("interned value removed prematurely");
+
+            if Arc::strong_count(arc) != 2 {
+                // Another thread has interned another copy
+                return;
+            }
+
+            shard.remove(&self.arc);
+
+            // Shrink the backing storage if the shard is less than 50% occupied.
+            if shard.len() * 2 < shard.capacity() {
+                shard.shrink_to_fit();
+            }
+        }
+    }
+}
+
+/// Compares interned `Ref`s using pointer equality.
+impl<T: Internable> PartialEq for Interned<T> {
+    #[inline]
+    fn eq(&self, other: &Self) -> bool {
+        Arc::ptr_eq(&self.arc, &other.arc)
+    }
+}
+
+impl<T: Internable> Eq for Interned<T> {}
+
+impl<T: Internable> AsRef<T> for Interned<T> {
+    #[inline]
+    fn as_ref(&self) -> &T {
+        &self.arc
+    }
+}
+
+impl<T: Internable> Deref for Interned<T> {
+    type Target = T;
+
+    #[inline]
+    fn deref(&self) -> &Self::Target {
+        &self.arc
+    }
+}
+
+impl<T: Internable> Clone for Interned<T> {
+    fn clone(&self) -> Self {
+        Self { arc: self.arc.clone() }
+    }
+}
+
+impl<T: Debug + Internable> Debug for Interned<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        (*self.arc).fmt(f)
+    }
+}
+
+pub struct InternStorage<T> {
+    map: OnceCell<InternMap<T>>,
+}
+
+impl<T> InternStorage<T> {
+    pub const fn new() -> Self {
+        Self { map: OnceCell::new() }
+    }
+}
+
+impl<T: Internable> InternStorage<T> {
+    fn get(&self) -> &InternMap<T> {
+        self.map.get_or_init(DashMap::default)
+    }
+}
+
+pub trait Internable: Hash + Eq + Sized + 'static {
+    fn storage() -> &'static InternStorage<Self>;
+}
+
+// region:`Internable` implementations
+
+macro_rules! impl_internable {
+    ( $($t:ty),+ $(,)? ) => { $(
+        impl Internable for $t {
+            fn storage() -> &'static InternStorage<Self> {
+                static STORAGE: InternStorage<$t> = InternStorage::new();
+                &STORAGE
+            }
+        }
+    )+ };
+}
+
+impl_internable!(crate::type_ref::TypeRef, crate::type_ref::TraitRef);
+
+// endregion