Merge #7557

7557: Intern `TypeRef`s in the containing `ItemTree` r=jonas-schievink a=jonas-schievink

This reduces the memory used by `ItemTreeQuery` by ~20%. As it turns out, `TypeRef` is very heavy, and is used a lot in `ItemTree`:

* Many types are frequently repeated throughout the same file. This is what this PR addresses by storing identical `TypeRef`s only once per `ItemTree`.
* The vast majority of `TypeRef` consist of a plain path with only a single element. "Type anchors" like in `<Ty>::func` are used incredibly rarely, and even multi-segment paths are much rarer than single-segment paths.

Co-authored-by: Jonas Schievink <[email protected]>

4 files changed, 67 insertions, 21 deletions

diff --git a/crates/hir_def/src/adt.rs b/crates/hir_def/src/adt.rs
index 06f0b9b18..ed36c3109 100644
--- a/crates/hir_def/src/adt.rs
+++ b/crates/hir_def/src/adt.rs
@@ -351,7 +351,7 @@ fn lower_field(
 ) -> FieldData {
     FieldData {
         name: field.name.clone(),
-        type_ref: field.type_ref.clone(),
+        type_ref: item_tree[field.type_ref].clone(),
         visibility: item_tree[override_visibility.unwrap_or(field.visibility)].clone(),
     }
 }
diff --git a/crates/hir_def/src/data.rs b/crates/hir_def/src/data.rs
index c2b0dc007..42fcca386 100644
--- a/crates/hir_def/src/data.rs
+++ b/crates/hir_def/src/data.rs
@@ -41,8 +41,8 @@ impl FunctionData {
 
         Arc::new(FunctionData {
             name: func.name.clone(),
-            params: func.params.to_vec(),
-            ret_type: func.ret_type.clone(),
+            params: func.params.iter().map(|id| item_tree[*id].clone()).collect(),
+            ret_type: item_tree[func.ret_type].clone(),
             attrs: item_tree.attrs(db, krate, ModItem::from(loc.id.value).into()).clone(),
             has_self_param: func.has_self_param,
             has_body: func.has_body,
@@ -75,7 +75,7 @@ impl TypeAliasData {
 
         Arc::new(TypeAliasData {
             name: typ.name.clone(),
-            type_ref: typ.type_ref.clone(),
+            type_ref: typ.type_ref.map(|id| item_tree[id].clone()),
             visibility: item_tree[typ.visibility].clone(),
             is_extern: typ.is_extern,
             bounds: typ.bounds.to_vec(),
@@ -144,8 +144,8 @@ impl ImplData {
 
         let item_tree = db.item_tree(impl_loc.id.file_id);
         let impl_def = &item_tree[impl_loc.id.value];
-        let target_trait = impl_def.target_trait.clone();
-        let target_type = impl_def.target_type.clone();
+        let target_trait = impl_def.target_trait.map(|id| item_tree[id].clone());
+        let target_type = item_tree[impl_def.target_type].clone();
         let is_negative = impl_def.is_negative;
         let module_id = impl_loc.container.module(db);
         let container = AssocContainerId::ImplId(id);
@@ -182,7 +182,7 @@ impl ConstData {
 
         Arc::new(ConstData {
             name: konst.name.clone(),
-            type_ref: konst.type_ref.clone(),
+            type_ref: item_tree[konst.type_ref].clone(),
             visibility: item_tree[konst.visibility].clone(),
         })
     }
@@ -205,7 +205,7 @@ impl StaticData {
 
         Arc::new(StaticData {
             name: Some(statik.name.clone()),
-            type_ref: statik.type_ref.clone(),
+            type_ref: item_tree[statik.type_ref].clone(),
             visibility: item_tree[statik.visibility].clone(),
             mutable: statik.mutable,
             is_extern: statik.is_extern,
diff --git a/crates/hir_def/src/item_tree.rs b/crates/hir_def/src/item_tree.rs
index 4bde67649..401556931 100644
--- a/crates/hir_def/src/item_tree.rs
+++ b/crates/hir_def/src/item_tree.rs
@@ -146,6 +146,7 @@ impl ItemTree {
                 macro_defs,
                 vis,
                 generics,
+                type_refs,
                 inner_items,
             } = &mut **data;
 
@@ -169,6 +170,8 @@ impl ItemTree {
 
             vis.arena.shrink_to_fit();
             generics.arena.shrink_to_fit();
+            type_refs.arena.shrink_to_fit();
+            type_refs.map.shrink_to_fit();
 
             inner_items.shrink_to_fit();
         }
@@ -279,6 +282,32 @@ static EMPTY_GENERICS: GenericParams = GenericParams {
     where_predicates: Vec::new(),
 };
 
+/// `TypeRef` interner.
+#[derive(Default, Debug, Eq, PartialEq)]
+struct TypeRefStorage {
+    arena: Arena<Arc<TypeRef>>,
+    map: FxHashMap<Arc<TypeRef>, Idx<Arc<TypeRef>>>,
+}
+
+impl TypeRefStorage {
+    // Note: We lie about the `Idx<TypeRef>` to hide the interner details.
+
+    fn intern(&mut self, ty: TypeRef) -> Idx<TypeRef> {
+        if let Some(id) = self.map.get(&ty) {
+            return Idx::from_raw(id.into_raw());
+        }
+
+        let ty = Arc::new(ty);
+        let idx = self.arena.alloc(ty.clone());
+        self.map.insert(ty, idx);
+        Idx::from_raw(idx.into_raw())
+    }
+
+    fn lookup(&self, id: Idx<TypeRef>) -> &TypeRef {
+        &self.arena[Idx::from_raw(id.into_raw())]
+    }
+}
+
 #[derive(Default, Debug, Eq, PartialEq)]
 struct ItemTreeData {
     imports: Arena<Import>,
@@ -301,6 +330,7 @@ struct ItemTreeData {
 
     vis: ItemVisibilities,
     generics: GenericParamsStorage,
+    type_refs: TypeRefStorage,
 
     inner_items: FxHashMap<FileAstId<ast::BlockExpr>, SmallVec<[ModItem; 1]>>,
 }
@@ -489,6 +519,14 @@ impl Index<GenericParamsId> for ItemTree {
     }
 }
 
+impl Index<Idx<TypeRef>> for ItemTree {
+    type Output = TypeRef;
+
+    fn index(&self, id: Idx<TypeRef>) -> &Self::Output {
+        self.data().type_refs.lookup(id)
+    }
+}
+
 impl<N: ItemTreeNode> Index<FileItemTreeId<N>> for ItemTree {
     type Output = N;
     fn index(&self, id: FileItemTreeId<N>) -> &N {
@@ -532,9 +570,9 @@ pub struct Function {
     /// Whether the function is located in an `extern` block (*not* whether it is an
     /// `extern "abi" fn`).
     pub is_extern: bool,
-    pub params: Box<[TypeRef]>,
+    pub params: Box<[Idx<TypeRef>]>,
     pub is_varargs: bool,
-    pub ret_type: TypeRef,
+    pub ret_type: Idx<TypeRef>,
     pub ast_id: FileAstId<ast::Fn>,
 }
 
@@ -581,7 +619,7 @@ pub struct Const {
     /// const _: () = ();
     pub name: Option<Name>,
     pub visibility: RawVisibilityId,
-    pub type_ref: TypeRef,
+    pub type_ref: Idx<TypeRef>,
     pub ast_id: FileAstId<ast::Const>,
 }
 
@@ -592,7 +630,7 @@ pub struct Static {
     pub mutable: bool,
     /// Whether the static is in an `extern` block.
     pub is_extern: bool,
-    pub type_ref: TypeRef,
+    pub type_ref: Idx<TypeRef>,
     pub ast_id: FileAstId<ast::Static>,
 }
 
@@ -609,8 +647,8 @@ pub struct Trait {
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct Impl {
     pub generic_params: GenericParamsId,
-    pub target_trait: Option<TypeRef>,
-    pub target_type: TypeRef,
+    pub target_trait: Option<Idx<TypeRef>>,
+    pub target_type: Idx<TypeRef>,
     pub is_negative: bool,
     pub items: Box<[AssocItem]>,
     pub ast_id: FileAstId<ast::Impl>,
@@ -623,7 +661,7 @@ pub struct TypeAlias {
     /// Bounds on the type alias itself. Only valid in trait declarations, eg. `type Assoc: Copy;`.
     pub bounds: Box<[TypeBound]>,
     pub generic_params: GenericParamsId,
-    pub type_ref: Option<TypeRef>,
+    pub type_ref: Option<Idx<TypeRef>>,
     pub is_extern: bool,
     pub ast_id: FileAstId<ast::TypeAlias>,
 }
@@ -806,6 +844,6 @@ pub enum Fields {
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Field {
     pub name: Name,
-    pub type_ref: TypeRef,
+    pub type_ref: Idx<TypeRef>,
     pub visibility: RawVisibilityId,
 }
diff --git a/crates/hir_def/src/item_tree/lower.rs b/crates/hir_def/src/item_tree/lower.rs
index de2177933..93cdca55d 100644
--- a/crates/hir_def/src/item_tree/lower.rs
+++ b/crates/hir_def/src/item_tree/lower.rs
@@ -364,6 +364,7 @@ impl Ctx {
                 params.push(type_ref);
             }
         }
+        let params = params.into_iter().map(|param| self.data().type_refs.intern(param)).collect();
 
         let mut is_varargs = false;
         if let Some(params) = func.param_list() {
@@ -385,6 +386,8 @@ impl Ctx {
             ret_type
         };
 
+        let ret_type = self.data().type_refs.intern(ret_type);
+
         let has_body = func.body().is_some();
 
         let ast_id = self.source_ast_id_map.ast_id(func);
@@ -396,7 +399,7 @@ impl Ctx {
             has_body,
             is_unsafe: func.unsafe_token().is_some(),
             is_extern: false,
-            params: params.into_boxed_slice(),
+            params,
             is_varargs,
             ret_type,
             ast_id,
@@ -657,6 +660,7 @@ impl Ctx {
                 generics.fill(&self.body_ctx, sm, node);
                 // lower `impl Trait` in arguments
                 for param in &*func.params {
+                    let param = self.data().type_refs.lookup(*param);
                     generics.fill_implicit_impl_trait_args(param);
                 }
             }
@@ -709,11 +713,15 @@ impl Ctx {
         self.data().vis.alloc(vis)
     }
 
-    fn lower_type_ref(&self, type_ref: &ast::Type) -> TypeRef {
-        TypeRef::from_ast(&self.body_ctx, type_ref.clone())
+    fn lower_type_ref(&mut self, type_ref: &ast::Type) -> Idx<TypeRef> {
+        let tyref = TypeRef::from_ast(&self.body_ctx, type_ref.clone());
+        self.data().type_refs.intern(tyref)
     }
-    fn lower_type_ref_opt(&self, type_ref: Option<ast::Type>) -> TypeRef {
-        type_ref.map(|ty| self.lower_type_ref(&ty)).unwrap_or(TypeRef::Error)
+    fn lower_type_ref_opt(&mut self, type_ref: Option<ast::Type>) -> Idx<TypeRef> {
+        match type_ref.map(|ty| self.lower_type_ref(&ty)) {
+            Some(it) => it,
+            None => self.data().type_refs.intern(TypeRef::Error),
+        }
     }
 
     /// Forces the visibility `vis` to be used for all items lowered during execution of `f`.