Merge #4947

4947: Replace `impls_in_trait` query with smarter use of `CrateImplDefs` r=matklad a=jonas-schievink

`impls_in_trait` was allocating a whopping ~400 MB of RAM when running analysis-stats on r-a itself.

Remove it, instead adding a query that computes a summary `CrateImplDefs` map for all transitive dependencies. This can probably still be made more efficient, but this already reduces the peak memory usage by 25% without much performance impact on analysis-stats.

**Before**:

```
Total: 34.962107188s, 2083mb allocated 2141mb resident
   422mb ImplsForTraitQuery (deps)
   250mb CrateDefMapQueryQuery
   147mb MacroArgQuery
   140mb TraitSolveQuery (deps)
    68mb InferQueryQuery (deps)
    62mb ImplDatumQuery (deps)
```

**After**:

```
Total: 35.261100358s, 1520mb allocated 1569mb resident
   250mb CrateDefMapQueryQuery
   147mb MacroArgQuery
   144mb TraitSolveQuery (deps)
    68mb InferQueryQuery (deps)
    61mb ImplDatumQuery (deps)
    45mb BodyQuery
    45mb ImplDatumQuery
```

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

4 files changed, 87 insertions, 56 deletions

diff --git a/crates/ra_hir_ty/src/db.rs b/crates/ra_hir_ty/src/db.rs
index bf71d38d6..7889b8d2c 100644
--- a/crates/ra_hir_ty/src/db.rs
+++ b/crates/ra_hir_ty/src/db.rs
@@ -3,15 +3,15 @@
 use std::sync::Arc;
 
 use hir_def::{
-    db::DefDatabase, DefWithBodyId, FunctionId, GenericDefId, ImplId, LocalFieldId, TraitId,
-    TypeParamId, VariantId,
+    db::DefDatabase, DefWithBodyId, FunctionId, GenericDefId, ImplId, LocalFieldId, TypeParamId,
+    VariantId,
 };
 use ra_arena::map::ArenaMap;
 use ra_db::{impl_intern_key, salsa, CrateId, Upcast};
 use ra_prof::profile;
 
 use crate::{
-    method_resolution::{CrateImplDefs, TyFingerprint},
+    method_resolution::CrateImplDefs,
     traits::{chalk, AssocTyValue, Impl},
     Binders, CallableDef, GenericPredicate, InferenceResult, OpaqueTyId, PolyFnSig,
     ReturnTypeImplTraits, Substs, TraitRef, Ty, TyDefId, TypeCtor, ValueTyDefId,
@@ -70,13 +70,8 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
     #[salsa::invoke(crate::method_resolution::CrateImplDefs::impls_in_crate_query)]
     fn impls_in_crate(&self, krate: CrateId) -> Arc<CrateImplDefs>;
 
-    #[salsa::invoke(crate::traits::impls_for_trait_query)]
-    fn impls_for_trait(
-        &self,
-        krate: CrateId,
-        trait_: TraitId,
-        self_ty_fp: Option<TyFingerprint>,
-    ) -> Arc<[ImplId]>;
+    #[salsa::invoke(crate::method_resolution::CrateImplDefs::impls_from_deps_query)]
+    fn impls_from_deps(&self, krate: CrateId) -> Arc<CrateImplDefs>;
 
     // Interned IDs for Chalk integration
     #[salsa::interned]
diff --git a/crates/ra_hir_ty/src/method_resolution.rs b/crates/ra_hir_ty/src/method_resolution.rs
index e83b39456..ed638c195 100644
--- a/crates/ra_hir_ty/src/method_resolution.rs
+++ b/crates/ra_hir_ty/src/method_resolution.rs
@@ -38,18 +38,53 @@ impl TyFingerprint {
     }
 }
 
+/// A queryable and mergeable collection of impls.
 #[derive(Debug, PartialEq, Eq)]
 pub struct CrateImplDefs {
-    impls: FxHashMap<TyFingerprint, Vec<ImplId>>,
+    inherent_impls: FxHashMap<TyFingerprint, Vec<ImplId>>,
     impls_by_trait: FxHashMap<TraitId, FxHashMap<Option<TyFingerprint>, Vec<ImplId>>>,
 }
 
 impl CrateImplDefs {
     pub(crate) fn impls_in_crate_query(db: &dyn HirDatabase, krate: CrateId) -> Arc<CrateImplDefs> {
         let _p = profile("impls_in_crate_query");
-        let mut res =
-            CrateImplDefs { impls: FxHashMap::default(), impls_by_trait: FxHashMap::default() };
+        let mut res = CrateImplDefs {
+            inherent_impls: FxHashMap::default(),
+            impls_by_trait: FxHashMap::default(),
+        };
+        res.fill(db, krate);
+
+        Arc::new(res)
+    }
+
+    /// Collects all impls from transitive dependencies of `krate` that may be used by `krate`.
+    ///
+    /// The full set of impls that can be used by `krate` is the returned map plus all the impls
+    /// from `krate` itself.
+    pub(crate) fn impls_from_deps_query(
+        db: &dyn HirDatabase,
+        krate: CrateId,
+    ) -> Arc<CrateImplDefs> {
+        let _p = profile("impls_from_deps_query");
+        let crate_graph = db.crate_graph();
+        let mut res = CrateImplDefs {
+            inherent_impls: FxHashMap::default(),
+            impls_by_trait: FxHashMap::default(),
+        };
 
+        // For each dependency, calculate `impls_from_deps` recursively, then add its own
+        // `impls_in_crate`.
+        // As we might visit crates multiple times, `merge` has to deduplicate impls to avoid
+        // wasting memory.
+        for dep in &crate_graph[krate].dependencies {
+            res.merge(&db.impls_from_deps(dep.crate_id));
+            res.merge(&db.impls_in_crate(dep.crate_id));
+        }
+
+        Arc::new(res)
+    }
+
+    fn fill(&mut self, db: &dyn HirDatabase, krate: CrateId) {
         let crate_def_map = db.crate_def_map(krate);
         for (_module_id, module_data) in crate_def_map.modules.iter() {
             for impl_id in module_data.scope.impls() {
@@ -57,7 +92,7 @@ impl CrateImplDefs {
                     Some(tr) => {
                         let self_ty = db.impl_self_ty(impl_id);
                         let self_ty_fp = TyFingerprint::for_impl(&self_ty.value);
-                        res.impls_by_trait
+                        self.impls_by_trait
                             .entry(tr.value.trait_)
                             .or_default()
                             .entry(self_ty_fp)
@@ -67,18 +102,36 @@ impl CrateImplDefs {
                     None => {
                         let self_ty = db.impl_self_ty(impl_id);
                         if let Some(self_ty_fp) = TyFingerprint::for_impl(&self_ty.value) {
-                            res.impls.entry(self_ty_fp).or_default().push(impl_id);
+                            self.inherent_impls.entry(self_ty_fp).or_default().push(impl_id);
                         }
                     }
                 }
             }
         }
+    }
 
-        Arc::new(res)
+    fn merge(&mut self, other: &Self) {
+        for (fp, impls) in &other.inherent_impls {
+            let vec = self.inherent_impls.entry(*fp).or_default();
+            vec.extend(impls);
+            vec.sort();
+            vec.dedup();
+        }
+
+        for (trait_, other_map) in &other.impls_by_trait {
+            let map = self.impls_by_trait.entry(*trait_).or_default();
+            for (fp, impls) in other_map {
+                let vec = map.entry(*fp).or_default();
+                vec.extend(impls);
+                vec.sort();
+                vec.dedup();
+            }
+        }
     }
+
     pub fn lookup_impl_defs(&self, ty: &Ty) -> impl Iterator<Item = ImplId> + '_ {
         let fingerprint = TyFingerprint::for_impl(ty);
-        fingerprint.and_then(|f| self.impls.get(&f)).into_iter().flatten().copied()
+        fingerprint.and_then(|f| self.inherent_impls.get(&f)).into_iter().flatten().copied()
     }
 
     pub fn lookup_impl_defs_for_trait(&self, tr: TraitId) -> impl Iterator<Item = ImplId> + '_ {
@@ -110,7 +163,7 @@ impl CrateImplDefs {
     }
 
     pub fn all_impls<'a>(&'a self) -> impl Iterator<Item = ImplId> + 'a {
-        self.impls
+        self.inherent_impls
             .values()
             .chain(self.impls_by_trait.values().flat_map(|m| m.values()))
             .flatten()
diff --git a/crates/ra_hir_ty/src/traits.rs b/crates/ra_hir_ty/src/traits.rs
index 892fbd6d1..6f43c3a22 100644
--- a/crates/ra_hir_ty/src/traits.rs
+++ b/crates/ra_hir_ty/src/traits.rs
@@ -7,9 +7,8 @@ use hir_def::{
 };
 use ra_db::{impl_intern_key, salsa, CrateId};
 use ra_prof::profile;
-use rustc_hash::FxHashSet;
 
-use crate::{db::HirDatabase, method_resolution::TyFingerprint, DebruijnIndex};
+use crate::{db::HirDatabase, DebruijnIndex};
 
 use super::{Canonical, GenericPredicate, HirDisplay, ProjectionTy, TraitRef, Ty, TypeWalk};
 
@@ -38,34 +37,6 @@ fn create_chalk_solver() -> chalk_solve::Solver<Interner> {
     solver_choice.into_solver()
 }
 
-/// Collects impls for the given trait in the whole dependency tree of `krate`.
-pub(crate) fn impls_for_trait_query(
-    db: &dyn HirDatabase,
-    krate: CrateId,
-    trait_: TraitId,
-    self_ty_fp: Option<TyFingerprint>,
-) -> Arc<[ImplId]> {
-    // FIXME: We could be a lot smarter here - because of the orphan rules and
-    // the fact that the trait and the self type need to be in the dependency
-    // tree of a crate somewhere for an impl to exist, we could skip looking in
-    // a lot of crates completely
-    let mut impls = FxHashSet::default();
-    // We call the query recursively here. On the one hand, this means we can
-    // reuse results from queries for different crates; on the other hand, this
-    // will only ever get called for a few crates near the root of the tree (the
-    // ones the user is editing), so this may actually be a waste of memory. I'm
-    // doing it like this mainly for simplicity for now.
-    for dep in &db.crate_graph()[krate].dependencies {
-        impls.extend(db.impls_for_trait(dep.crate_id, trait_, self_ty_fp).iter());
-    }
-    let crate_impl_defs = db.impls_in_crate(krate);
-    match self_ty_fp {
-        Some(fp) => impls.extend(crate_impl_defs.lookup_impl_defs_for_trait_and_ty(trait_, fp)),
-        None => impls.extend(crate_impl_defs.lookup_impl_defs_for_trait(trait_)),
-    }
-    impls.into_iter().collect()
-}
-
 /// A set of clauses that we assume to be true. E.g. if we are inside this function:
 /// ```rust
 /// fn foo<T: Default>(t: T) {}
diff --git a/crates/ra_hir_ty/src/traits/chalk.rs b/crates/ra_hir_ty/src/traits/chalk.rs
index a72a82f5a..2f35d6d49 100644
--- a/crates/ra_hir_ty/src/traits/chalk.rs
+++ b/crates/ra_hir_ty/src/traits/chalk.rs
@@ -74,14 +74,26 @@ impl<'a> chalk_solve::RustIrDatabase<Interner> for ChalkContext<'a> {
         // Note: Since we're using impls_for_trait, only impls where the trait
         // can be resolved should ever reach Chalk. `impl_datum` relies on that
         // and will panic if the trait can't be resolved.
-        let mut result: Vec<_> = self
-            .db
-            .impls_for_trait(self.krate, trait_, self_ty_fp)
-            .iter()
-            .copied()
-            .map(Impl::ImplDef)
-            .map(|impl_| impl_.to_chalk(self.db))
-            .collect();
+        let in_deps = self.db.impls_from_deps(self.krate);
+        let in_self = self.db.impls_in_crate(self.krate);
+        let impl_maps = [in_deps, in_self];
+
+        let id_to_chalk = |id: hir_def::ImplId| Impl::ImplDef(id).to_chalk(self.db);
+
+        let mut result: Vec<_> = match self_ty_fp {
+            Some(fp) => impl_maps
+                .iter()
+                .flat_map(|crate_impl_defs| {
+                    crate_impl_defs.lookup_impl_defs_for_trait_and_ty(trait_, fp).map(id_to_chalk)
+                })
+                .collect(),
+            None => impl_maps
+                .iter()
+                .flat_map(|crate_impl_defs| {
+                    crate_impl_defs.lookup_impl_defs_for_trait(trait_).map(id_to_chalk)
+                })
+                .collect(),
+        };
 
         let arg: Option<Ty> =
             parameters.get(1).map(|p| from_chalk(self.db, p.assert_ty_ref(&Interner).clone()));