Look up impls by self type

This speeds up inference in analysis-stats by ~30% (even more with the recursive
solver).

4 files changed, 65 insertions, 14 deletions

diff --git a/crates/ra_hir_ty/src/db.rs b/crates/ra_hir_ty/src/db.rs
index 1462b053f..33da16b48 100644
--- a/crates/ra_hir_ty/src/db.rs
+++ b/crates/ra_hir_ty/src/db.rs
@@ -11,7 +11,7 @@ use ra_db::{impl_intern_key, salsa, CrateId, Upcast};
 use ra_prof::profile;
 
 use crate::{
-    method_resolution::CrateImplDefs,
+    method_resolution::{CrateImplDefs, TyFingerprint},
     traits::{chalk, AssocTyValue, Impl},
     Binders, CallableDef, GenericPredicate, InferenceResult, PolyFnSig, Substs, TraitRef, Ty,
     TyDefId, TypeCtor, ValueTyDefId,
@@ -65,7 +65,12 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
     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) -> Arc<[ImplId]>;
+    fn impls_for_trait(
+        &self,
+        krate: CrateId,
+        trait_: TraitId,
+        self_ty_fp: Option<TyFingerprint>,
+    ) -> Arc<[ImplId]>;
 
     // 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 74a0bc7db..657284fd0 100644
--- a/crates/ra_hir_ty/src/method_resolution.rs
+++ b/crates/ra_hir_ty/src/method_resolution.rs
@@ -34,7 +34,7 @@ impl TyFingerprint {
     /// Creates a TyFingerprint for looking up an impl. Only certain types can
     /// have impls: if we have some `struct S`, we can have an `impl S`, but not
     /// `impl &S`. Hence, this will return `None` for reference types and such.
-    fn for_impl(ty: &Ty) -> Option<TyFingerprint> {
+    pub(crate) fn for_impl(ty: &Ty) -> Option<TyFingerprint> {
         match ty {
             Ty::Apply(a_ty) => Some(TyFingerprint::Apply(a_ty.ctor)),
             _ => None,
@@ -45,7 +45,7 @@ impl TyFingerprint {
 #[derive(Debug, PartialEq, Eq)]
 pub struct CrateImplDefs {
     impls: FxHashMap<TyFingerprint, Vec<ImplId>>,
-    impls_by_trait: FxHashMap<TraitId, Vec<ImplId>>,
+    impls_by_trait: FxHashMap<TraitId, FxHashMap<Option<TyFingerprint>, Vec<ImplId>>>,
 }
 
 impl CrateImplDefs {
@@ -59,7 +59,14 @@ impl CrateImplDefs {
             for impl_id in module_data.scope.impls() {
                 match db.impl_trait(impl_id) {
                     Some(tr) => {
-                        res.impls_by_trait.entry(tr.value.trait_).or_default().push(impl_id);
+                        let self_ty = db.impl_self_ty(impl_id);
+                        let self_ty_fp = TyFingerprint::for_impl(&self_ty.value);
+                        res.impls_by_trait
+                            .entry(tr.value.trait_)
+                            .or_default()
+                            .entry(self_ty_fp)
+                            .or_default()
+                            .push(impl_id);
                     }
                     None => {
                         let self_ty = db.impl_self_ty(impl_id);
@@ -79,11 +86,39 @@ impl CrateImplDefs {
     }
 
     pub fn lookup_impl_defs_for_trait(&self, tr: TraitId) -> impl Iterator<Item = ImplId> + '_ {
-        self.impls_by_trait.get(&tr).into_iter().flatten().copied()
+        self.impls_by_trait
+            .get(&tr)
+            .into_iter()
+            .flat_map(|m| m.values().flat_map(|v| v.iter().copied()))
+    }
+
+    pub fn lookup_impl_defs_for_trait_and_ty(
+        &self,
+        tr: TraitId,
+        fp: TyFingerprint,
+    ) -> impl Iterator<Item = ImplId> + '_ {
+        self.impls_by_trait
+            .get(&tr)
+            .and_then(|m| m.get(&Some(fp)))
+            .into_iter()
+            .flatten()
+            .copied()
+            .chain(
+                self.impls_by_trait
+                    .get(&tr)
+                    .and_then(|m| m.get(&None))
+                    .into_iter()
+                    .flatten()
+                    .copied(),
+            )
     }
 
     pub fn all_impls<'a>(&'a self) -> impl Iterator<Item = ImplId> + 'a {
-        self.impls.values().chain(self.impls_by_trait.values()).flatten().copied()
+        self.impls
+            .values()
+            .chain(self.impls_by_trait.values().flat_map(|m| m.values()))
+            .flatten()
+            .copied()
     }
 }
 
diff --git a/crates/ra_hir_ty/src/traits.rs b/crates/ra_hir_ty/src/traits.rs
index 21e233379..43d8d1e80 100644
--- a/crates/ra_hir_ty/src/traits.rs
+++ b/crates/ra_hir_ty/src/traits.rs
@@ -7,7 +7,7 @@ use ra_db::{impl_intern_key, salsa, CrateId};
 use ra_prof::profile;
 use rustc_hash::FxHashSet;
 
-use crate::{db::HirDatabase, DebruijnIndex};
+use crate::{db::HirDatabase, method_resolution::TyFingerprint, DebruijnIndex};
 
 use super::{Canonical, GenericPredicate, HirDisplay, ProjectionTy, TraitRef, Ty, TypeWalk};
 
@@ -40,7 +40,12 @@ 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
@@ -48,10 +53,13 @@ pub(crate) fn impls_for_trait_query(
     // 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_).iter());
+        impls.extend(db.impls_for_trait(dep.crate_id, trait_, self_ty_fp).iter());
     }
     let crate_impl_defs = db.impls_in_crate(krate);
-    impls.extend(crate_impl_defs.lookup_impl_defs_for_trait(trait_));
+    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()
 }
 
diff --git a/crates/ra_hir_ty/src/traits/chalk.rs b/crates/ra_hir_ty/src/traits/chalk.rs
index c5f1b5232..e05fea843 100644
--- a/crates/ra_hir_ty/src/traits/chalk.rs
+++ b/crates/ra_hir_ty/src/traits/chalk.rs
@@ -16,8 +16,8 @@ use ra_db::{
 
 use super::{builtin, AssocTyValue, Canonical, ChalkContext, Impl, Obligation};
 use crate::{
-    db::HirDatabase, display::HirDisplay, utils::generics, ApplicationTy, GenericPredicate,
-    ProjectionTy, Substs, TraitRef, Ty, TypeCtor,
+    db::HirDatabase, display::HirDisplay, method_resolution::TyFingerprint, utils::generics,
+    ApplicationTy, GenericPredicate, ProjectionTy, Substs, TraitRef, Ty, TypeCtor,
 };
 
 pub(super) mod tls;
@@ -647,19 +647,22 @@ impl<'a> chalk_solve::RustIrDatabase<Interner> for ChalkContext<'a> {
         debug!("impls_for_trait {:?}", trait_id);
         let trait_: hir_def::TraitId = from_chalk(self.db, trait_id);
 
+        let ty: Ty = from_chalk(self.db, parameters[0].assert_ty_ref(&Interner).clone());
+
+        let self_ty_fp = TyFingerprint::for_impl(&ty);
+
         // 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_)
+            .impls_for_trait(self.krate, trait_, self_ty_fp)
             .iter()
             .copied()
             .map(Impl::ImplDef)
             .map(|impl_| impl_.to_chalk(self.db))
             .collect();
 
-        let ty: Ty = from_chalk(self.db, parameters[0].assert_ty_ref(&Interner).clone());
         let arg: Option<Ty> =
             parameters.get(1).map(|p| from_chalk(self.db, p.assert_ty_ref(&Interner).clone()));