Auto merge of rust-lang#14716 - Veykril:trait-deps-query, r=Veykril

Don't merge trait_impls_in_deps results

Closes rust-lang/rust-analyzer#9167

We effectively keep a lot of copies of `TraitImpl` results in memory by merging them into the result of this query. This wastes a lot of memory unnecessarily. The change here brings memory usage of self down by 50mb (at a slight cost of trait solving unfortunately, though I'm hopeful it was just noice)

crates/hir-ty/src/chalk_db.rs

@@ -1,5 +1,6 @@
 //! The implementation of `RustIrDatabase` for Chalk, which provides information
 //! about the code that Chalk needs.
+use core::ops;
 use std::{iter, sync::Arc};
 
 use tracing::debug;
@@ -126,7 +127,6 @@ impl<'a> chalk_solve::RustIrDatabase<Interner> for ChalkContext<'a> {
         let in_deps = self.db.trait_impls_in_deps(self.krate);
         let in_self = self.db.trait_impls_in_crate(self.krate);
 
-        let impl_maps = [in_deps, in_self];
         let block_impls = iter::successors(self.block, |&block_id| {
             cov_mark::hit!(block_local_impls);
             self.db.block_def_map(block_id).parent().and_then(|module| module.containing_block())
@@ -146,29 +146,31 @@ impl<'a> chalk_solve::RustIrDatabase<Interner> for ChalkContext<'a> {
         match fps {
             [] => {
                 debug!("Unrestricted search for {:?} impls...", trait_);
-                let mut f = |impls: Arc<TraitImpls>| {
+                let mut f = |impls: &TraitImpls| {
                     result.extend(impls.for_trait(trait_).map(id_to_chalk));
                 };
-                impl_maps.into_iter().chain(block_impls).for_each(&mut f);
+                f(&in_self);
+                in_deps.iter().map(ops::Deref::deref).for_each(&mut f);
+                block_impls.for_each(|it| f(&it));
                 def_blocks
                     .into_iter()
                     .flatten()
-                    .map(|it| self.db.trait_impls_in_block(it))
-                    .for_each(f);
+                    .for_each(|it| f(&self.db.trait_impls_in_block(it)));
             }
             fps => {
                 let mut f =
-                    |impls: Arc<TraitImpls>| {
+                    |impls: &TraitImpls| {
                         result.extend(fps.iter().flat_map(|fp| {
                             impls.for_trait_and_self_ty(trait_, *fp).map(id_to_chalk)
                         }));
                     };
-                impl_maps.into_iter().chain(block_impls).for_each(&mut f);
+                f(&in_self);
+                in_deps.iter().map(ops::Deref::deref).for_each(&mut f);
+                block_impls.for_each(|it| f(&it));
                 def_blocks
                     .into_iter()
                     .flatten()
-                    .map(|it| self.db.trait_impls_in_block(it))
-                    .for_each(f);
+                    .for_each(|it| f(&self.db.trait_impls_in_block(it)));
             }
         }
 

crates/hir-ty/src/db.rs

@@ -134,7 +134,7 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
     fn trait_impls_in_block(&self, block: BlockId) -> Arc<TraitImpls>;
 
     #[salsa::invoke(TraitImpls::trait_impls_in_deps_query)]
-    fn trait_impls_in_deps(&self, krate: CrateId) -> Arc<TraitImpls>;
+    fn trait_impls_in_deps(&self, krate: CrateId) -> Arc<[Arc<TraitImpls>]>;
 
     // Interned IDs for Chalk integration
     #[salsa::interned]

crates/hir-ty/src/method_resolution.rs

@@ -160,17 +160,13 @@ impl TraitImpls {
         Arc::new(impls)
     }
 
-    pub(crate) fn trait_impls_in_deps_query(db: &dyn HirDatabase, krate: CrateId) -> Arc<Self> {
+    pub(crate) fn trait_impls_in_deps_query(
+        db: &dyn HirDatabase,
+        krate: CrateId,
+    ) -> Arc<[Arc<Self>]> {
         let _p = profile::span("trait_impls_in_deps_query").detail(|| format!("{krate:?}"));
         let crate_graph = db.crate_graph();
-        let mut res = Self { map: FxHashMap::default() };
-
-        for krate in crate_graph.transitive_deps(krate) {
-            res.merge(&db.trait_impls_in_crate(krate));
-        }
-        res.shrink_to_fit();
-
-        Arc::new(res)
+        crate_graph.transitive_deps(krate).map(|krate| db.trait_impls_in_crate(krate)).collect()
     }
 
     fn shrink_to_fit(&mut self) {
@@ -209,15 +205,6 @@ impl TraitImpls {
         }
     }
 
-    fn merge(&mut self, other: &Self) {
-        for (trait_, other_map) in &other.map {
-            let map = self.map.entry(*trait_).or_default();
-            for (fp, impls) in other_map {
-                map.entry(*fp).or_default().extend(impls);
-            }
-        }
-    }
-
     /// Queries all trait impls for the given type.
     pub fn for_self_ty_without_blanket_impls(
         &self,
@@ -713,10 +700,12 @@ fn lookup_impl_assoc_item_for_trait_ref(
     env: Arc<TraitEnvironment>,
     name: &Name,
 ) -> Option<(AssocItemId, Substitution)> {
+    let hir_trait_id = trait_ref.hir_trait_id();
     let self_ty = trait_ref.self_type_parameter(Interner);
     let self_ty_fp = TyFingerprint::for_trait_impl(&self_ty)?;
     let impls = db.trait_impls_in_deps(env.krate);
-    let impls = impls.for_trait_and_self_ty(trait_ref.hir_trait_id(), self_ty_fp);
+    let impls =
+        impls.iter().flat_map(|impls| impls.for_trait_and_self_ty(hir_trait_id, self_ty_fp));
 
     let table = InferenceTable::new(db, env);