Auto merge of #16555 - davidbarsky:david/speedup-completions-by-exploiting-orphan-rules, r=Veykril

performance: Speed up Method Completions By Taking Advantage of Orphan Rules

(Continues https://github.com/rust-lang/rust-analyzer/pull/16498)

This PR speeds up method completions by doing two things without regressing `analysis-stats`[^1]:
- Filter candidate traits prior to calling `iterate_path_candidates` by relying on orphan rules (see below for a slightly more in-depth explanation). When generating completions [on `slog::Logger`](https://github.com/oxidecomputer/omicron/blob/5e9e59c312cd787ba50cf6776f220951917dfa14/common/src/ledger.rs#L78) in `oxidecomputer/omicron` as a test, this PR halved my completion times—it's now 454ms cold and 281ms warm. Before this PR, it was 808ms cold and 579ms warm.
- Inline some of the method candidate checks into `is_valid_method_candidate` and remove some unnecessary visibility checks. This was suggested by `@Veykril` in [this comment](https://github.com/rust-lang/rust-analyzer/pull/16498#issuecomment-1929864427).

We filter candidate traits by taking advantage of orphan rules. For additional details, I'll rely on `@WaffleLapkin's` explanation  [from Zulip](https://rust-lang.zulipchat.com/#narrow/stream/185405-t-compiler.2Frust-analyzer/topic/Trait.20Checking/near/420942417):

> A type `A` can only implements traits which
> 1. Have a blanket implementation (`impl<T> Trait for T {}`)
> 2. Have implementation for `A` (`impl Trait for A {}`)
>
> Blanket implementation can only exist in `Trait`'s crate. Implementation for `A` can only exist in `A`'s or `Trait`'s crate.

Big thanks to Waffle for its keen observation!

---

I think some additional improvements are possible:
- `for_trait_and_self_ty` seemingly does not distinguish between `&T`, `&mut T`, or `T`, resulting in seemingly irrelevant traits like `tokio::io::AsyncWrite` being being included for, e.g., `&slog::Logger`. I don't know they're being considered due to the [autoref/autoderef behavior](https://github.com/rust-lang/rust-analyzer/blob/a02a219773629686bd8ff123ca1aa995fa50d976/crates/hir-ty/src/method_resolution.rs#L945-L962), but I wonder if it'd make sense to filter by mutability earlier and not consider trait implementations that require `&mut T` when we only have a `&T`.
- The method completions [spend a _lot_ of time in unification](https://rust-lang.zulipchat.com/#narrow/stream/185405-t-compiler.2Frust-analyzer/topic/Trait.20Checking/near/421072356), and while there might be low-hanging fruit there, it might make more sense to wait for the new trait solver in `rustc`. I dunno.

[^1]: The filtering occurs outside of typechecking, after all.

6 files changed, 243 insertions, 4 deletions

diff --git a/crates/hir-ty/src/infer/unify.rs b/crates/hir-ty/src/infer/unify.rs
index 709760b64fd..4a181f4a325 100644
--- a/crates/hir-ty/src/infer/unify.rs
+++ b/crates/hir-ty/src/infer/unify.rs
@@ -457,6 +457,7 @@ impl<'a> InferenceTable<'a> {
     }
 
     /// Unify two relatable values (e.g. `Ty`) and register new trait goals that arise from that.
+    #[tracing::instrument(skip_all)]
     pub(crate) fn unify<T: ?Sized + Zip<Interner>>(&mut self, ty1: &T, ty2: &T) -> bool {
         let result = match self.try_unify(ty1, ty2) {
             Ok(r) => r,
diff --git a/crates/hir-ty/src/method_resolution.rs b/crates/hir-ty/src/method_resolution.rs
index a4baf572d9e..4bb412d01c2 100644
--- a/crates/hir-ty/src/method_resolution.rs
+++ b/crates/hir-ty/src/method_resolution.rs
@@ -254,6 +254,11 @@ impl TraitImpls {
             .flat_map(|v| v.iter().copied())
     }
 
+    /// Queries whether `self_ty` has potentially applicable implementations of `trait_`.
+    pub fn has_impls_for_trait_and_self_ty(&self, trait_: TraitId, self_ty: TyFingerprint) -> bool {
+        self.for_trait_and_self_ty(trait_, self_ty).next().is_some()
+    }
+
     pub fn all_impls(&self) -> impl Iterator<Item = ImplId> + '_ {
         self.map.values().flat_map(|map| map.values().flat_map(|v| v.iter().copied()))
     }
@@ -1170,7 +1175,7 @@ fn iterate_trait_method_candidates(
         for &(_, item) in data.items.iter() {
             // Don't pass a `visible_from_module` down to `is_valid_candidate`,
             // since only inherent methods should be included into visibility checking.
-            let visible = match is_valid_candidate(table, name, receiver_ty, item, self_ty, None) {
+            let visible = match is_valid_method_candidate(table, name, receiver_ty, item, self_ty) {
                 IsValidCandidate::Yes => true,
                 IsValidCandidate::NotVisible => false,
                 IsValidCandidate::No => continue,
@@ -1414,6 +1419,74 @@ fn is_valid_candidate(
     }
 }
 
+/// Checks whether a given `AssocItemId` is applicable for `receiver_ty`.
+///
+/// This method should *only* be called by [`iterate_trait_method_candidates`],
+/// as it is responsible for determining applicability in completions.
+#[tracing::instrument(skip_all, fields(name))]
+fn is_valid_method_candidate(
+    table: &mut InferenceTable<'_>,
+    name: Option<&Name>,
+    receiver_ty: Option<&Ty>,
+    item: AssocItemId,
+    self_ty: &Ty,
+) -> IsValidCandidate {
+    let db = table.db;
+    match item {
+        AssocItemId::FunctionId(fn_id) => {
+            let data = db.function_data(fn_id);
+
+            check_that!(name.map_or(true, |n| n == &data.name));
+
+            table.run_in_snapshot(|table| {
+                let container = fn_id.lookup(db.upcast()).container;
+                let (impl_subst, expect_self_ty) = match container {
+                    ItemContainerId::ImplId(it) => {
+                        let subst = TyBuilder::subst_for_def(db, it, None)
+                            .fill_with_inference_vars(table)
+                            .build();
+                        let self_ty = db.impl_self_ty(it).substitute(Interner, &subst);
+                        (subst, self_ty)
+                    }
+                    ItemContainerId::TraitId(it) => {
+                        let subst = TyBuilder::subst_for_def(db, it, None)
+                            .fill_with_inference_vars(table)
+                            .build();
+                        let self_ty = subst.at(Interner, 0).assert_ty_ref(Interner).clone();
+                        (subst, self_ty)
+                    }
+                    _ => unreachable!(),
+                };
+
+                check_that!(table.unify(&expect_self_ty, self_ty));
+
+                if let Some(receiver_ty) = receiver_ty {
+                    check_that!(data.has_self_param());
+
+                    let fn_subst = TyBuilder::subst_for_def(db, fn_id, Some(impl_subst.clone()))
+                        .fill_with_inference_vars(table)
+                        .build();
+
+                    let sig = db.callable_item_signature(fn_id.into());
+                    let expected_receiver =
+                        sig.map(|s| s.params()[0].clone()).substitute(Interner, &fn_subst);
+
+                    check_that!(table.unify(receiver_ty, &expected_receiver));
+                }
+
+                IsValidCandidate::Yes
+            })
+        }
+        AssocItemId::ConstId(c) => {
+            check_that!(receiver_ty.is_none());
+            check_that!(name.map_or(true, |n| db.const_data(c).name.as_ref() == Some(n)));
+
+            IsValidCandidate::Yes
+        }
+        _ => IsValidCandidate::No,
+    }
+}
+
 enum IsValidCandidate {
     Yes,
     No,
@@ -1441,6 +1514,8 @@ fn is_valid_fn_candidate(
     }
     table.run_in_snapshot(|table| {
         let container = fn_id.lookup(db.upcast()).container;
+
+        let _p = tracing::span!(tracing::Level::INFO, "subst_for_def").entered();
         let (impl_subst, expect_self_ty) = match container {
             ItemContainerId::ImplId(it) => {
                 let subst =
@@ -1460,6 +1535,7 @@ fn is_valid_fn_candidate(
         check_that!(table.unify(&expect_self_ty, self_ty));
 
         if let Some(receiver_ty) = receiver_ty {
+            let _p = tracing::span!(tracing::Level::INFO, "check_receiver_ty").entered();
             check_that!(data.has_self_param());
 
             let fn_subst = TyBuilder::subst_for_def(db, fn_id, Some(impl_subst.clone()))
@@ -1474,6 +1550,7 @@ fn is_valid_fn_candidate(
         }
 
         if let ItemContainerId::ImplId(impl_id) = container {
+            let _p = tracing::span!(tracing::Level::INFO, "check_item_container").entered();
             // We need to consider the bounds on the impl to distinguish functions of the same name
             // for a type.
             let predicates = db.generic_predicates(impl_id.into());
diff --git a/crates/hir-ty/src/traits.rs b/crates/hir-ty/src/traits.rs
index b2232b920aa..71a4ab020b3 100644
--- a/crates/hir-ty/src/traits.rs
+++ b/crates/hir-ty/src/traits.rs
@@ -139,6 +139,7 @@ fn solve(
     block: Option<BlockId>,
     goal: &chalk_ir::UCanonical<chalk_ir::InEnvironment<chalk_ir::Goal<Interner>>>,
 ) -> Option<chalk_solve::Solution<Interner>> {
+    let _p = tracing::span!(tracing::Level::INFO, "solve", ?krate, ?block).entered();
     let context = ChalkContext { db, krate, block };
     tracing::debug!("solve goal: {:?}", goal);
     let mut solver = create_chalk_solver();
diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index a0237e3b90b..3147b591413 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -4239,6 +4239,10 @@ impl Type {
         }
     }
 
+    pub fn fingerprint_for_trait_impl(&self) -> Option<TyFingerprint> {
+        TyFingerprint::for_trait_impl(&self.ty)
+    }
+
     pub(crate) fn canonical(&self) -> Canonical<Ty> {
         hir_ty::replace_errors_with_variables(&self.ty)
     }
diff --git a/crates/ide-completion/src/tests/flyimport.rs b/crates/ide-completion/src/tests/flyimport.rs
index fff193ba4c9..d2227d23cd7 100644
--- a/crates/ide-completion/src/tests/flyimport.rs
+++ b/crates/ide-completion/src/tests/flyimport.rs
@@ -375,6 +375,135 @@ fn main() {
 }
 
 #[test]
+fn trait_method_fuzzy_completion_aware_of_fundamental_boxes() {
+    let fixture = r#"
+//- /fundamental.rs crate:fundamental
+#[lang = "owned_box"]
+#[fundamental]
+pub struct Box<T>(T);
+//- /foo.rs crate:foo
+pub trait TestTrait {
+    fn some_method(&self);
+}
+//- /main.rs crate:main deps:foo,fundamental
+struct TestStruct;
+
+impl foo::TestTrait for fundamental::Box<TestStruct> {
+    fn some_method(&self) {}
+}
+
+fn main() {
+    let t = fundamental::Box(TestStruct);
+    t.$0
+}
+"#;
+
+    check(
+        fixture,
+        expect![[r#"
+        me some_method() (use foo::TestTrait) fn(&self)
+    "#]],
+    );
+
+    check_edit(
+        "some_method",
+        fixture,
+        r#"
+use foo::TestTrait;
+
+struct TestStruct;
+
+impl foo::TestTrait for fundamental::Box<TestStruct> {
+    fn some_method(&self) {}
+}
+
+fn main() {
+    let t = fundamental::Box(TestStruct);
+    t.some_method()$0
+}
+"#,
+    );
+}
+
+#[test]
+fn trait_method_fuzzy_completion_aware_of_fundamental_references() {
+    let fixture = r#"
+//- /foo.rs crate:foo
+pub trait TestTrait {
+    fn some_method(&self);
+}
+//- /main.rs crate:main deps:foo
+struct TestStruct;
+
+impl foo::TestTrait for &TestStruct {
+    fn some_method(&self) {}
+}
+
+fn main() {
+    let t = &TestStruct;
+    t.$0
+}
+"#;
+
+    check(
+        fixture,
+        expect![[r#"
+        me some_method() (use foo::TestTrait) fn(&self)
+    "#]],
+    );
+
+    check_edit(
+        "some_method",
+        fixture,
+        r#"
+use foo::TestTrait;
+
+struct TestStruct;
+
+impl foo::TestTrait for &TestStruct {
+    fn some_method(&self) {}
+}
+
+fn main() {
+    let t = &TestStruct;
+    t.some_method()$0
+}
+"#,
+    );
+}
+
+#[test]
+fn trait_method_fuzzy_completion_aware_of_unit_type() {
+    let fixture = r#"
+//- /test_trait.rs crate:test_trait
+pub trait TestInto<T> {
+    fn into(self) -> T;
+}
+
+//- /main.rs crate:main deps:test_trait
+struct A;
+
+impl test_trait::TestInto<A> for () {
+    fn into(self) -> A {
+        A
+    }
+}
+
+fn main() {
+    let a = ();
+    a.$0
+}
+"#;
+
+    check(
+        fixture,
+        expect![[r#"
+    me into() (use test_trait::TestInto) fn(self) -> T
+    "#]],
+    );
+}
+
+#[test]
 fn trait_method_from_alias() {
     let fixture = r#"
 //- /lib.rs crate:dep
diff --git a/crates/ide-db/src/imports/import_assets.rs b/crates/ide-db/src/imports/import_assets.rs
index a71d8e9002d..c597555a3bf 100644
--- a/crates/ide-db/src/imports/import_assets.rs
+++ b/crates/ide-db/src/imports/import_assets.rs
@@ -1,8 +1,9 @@
 //! Look up accessible paths for items.
 
 use hir::{
-    AsAssocItem, AssocItem, AssocItemContainer, Crate, ItemInNs, ModPath, Module, ModuleDef, Name,
-    PathResolution, PrefixKind, ScopeDef, Semantics, SemanticsScope, Type,
+    db::HirDatabase, AsAssocItem, AssocItem, AssocItemContainer, Crate, HasCrate, ItemInNs,
+    ModPath, Module, ModuleDef, Name, PathResolution, PrefixKind, ScopeDef, Semantics,
+    SemanticsScope, Trait, Type,
 };
 use itertools::{EitherOrBoth, Itertools};
 use rustc_hash::{FxHashMap, FxHashSet};
@@ -517,7 +518,7 @@ fn trait_applicable_items(
     let related_traits = inherent_traits.chain(env_traits).collect::<FxHashSet<_>>();
 
     let mut required_assoc_items = FxHashSet::default();
-    let trait_candidates: FxHashSet<_> = items_locator::items_with_name(
+    let mut trait_candidates: FxHashSet<_> = items_locator::items_with_name(
         sema,
         current_crate,
         trait_candidate.assoc_item_name.clone(),
@@ -538,6 +539,32 @@ fn trait_applicable_items(
     })
     .collect();
 
+    trait_candidates.retain(|&candidate_trait_id| {
+        // we care about the following cases:
+        // 1. Trait's definition crate
+        // 2. Definition crates for all trait's generic arguments
+        //     a. This is recursive for fundamental types: `Into<Box<A>> for ()`` is OK, but
+        //        `Into<Vec<A>> for ()`` is *not*.
+        // 3. Receiver type definition crate
+        //    a. This is recursive for fundamental types
+        let defining_crate_for_trait = Trait::from(candidate_trait_id).krate(db);
+        let Some(receiver) = trait_candidate.receiver_ty.fingerprint_for_trait_impl() else {
+            return false;
+        };
+        let definitions_exist_in_trait_crate = db
+            .trait_impls_in_crate(defining_crate_for_trait.into())
+            .has_impls_for_trait_and_self_ty(candidate_trait_id, receiver);
+
+        // this is a closure for laziness: if `definitions_exist_in_trait_crate` is true,
+        // we can avoid a second db lookup.
+        let definitions_exist_in_receiver_crate = || {
+            db.trait_impls_in_crate(trait_candidate.receiver_ty.krate(db).into())
+                .has_impls_for_trait_and_self_ty(candidate_trait_id, receiver)
+        };
+
+        definitions_exist_in_trait_crate || definitions_exist_in_receiver_crate()
+    });
+
     let mut located_imports = FxHashSet::default();
     let mut trait_import_paths = FxHashMap::default();