Make block-local trait impls work

As long as either the trait or the implementing type are defined in the
same block.

1 files changed, 38 insertions, 26 deletions

diff --git a/crates/hir_ty/src/method_resolution.rs b/crates/hir_ty/src/method_resolution.rs
index f3d390961..3d233b1e2 100644
--- a/crates/hir_ty/src/method_resolution.rs
+++ b/crates/hir_ty/src/method_resolution.rs
@@ -8,7 +8,7 @@ use arrayvec::ArrayVec;
 use base_db::{CrateId, Edition};
 use chalk_ir::{cast::Cast, Mutability, UniverseIndex};
 use hir_def::{
-    lang_item::LangItemTarget, nameres::DefMap, AssocContainerId, AssocItemId, FunctionId,
+    lang_item::LangItemTarget, nameres::DefMap, AssocContainerId, AssocItemId, BlockId, FunctionId,
     GenericDefId, HasModule, ImplId, Lookup, ModuleId, TraitId,
 };
 use hir_expand::name::Name;
@@ -139,35 +139,47 @@ impl TraitImpls {
         let mut impls = Self { map: FxHashMap::default() };
 
         let crate_def_map = db.crate_def_map(krate);
-        collect_def_map(db, &crate_def_map, &mut impls);
+        impls.collect_def_map(db, &crate_def_map);
 
         return Arc::new(impls);
+    }
 
-        fn collect_def_map(db: &dyn HirDatabase, def_map: &DefMap, impls: &mut TraitImpls) {
-            for (_module_id, module_data) in def_map.modules() {
-                for impl_id in module_data.scope.impls() {
-                    let target_trait = match db.impl_trait(impl_id) {
-                        Some(tr) => tr.skip_binders().hir_trait_id(),
-                        None => continue,
-                    };
-                    let self_ty = db.impl_self_ty(impl_id);
-                    let self_ty_fp = TyFingerprint::for_trait_impl(self_ty.skip_binders());
-                    impls
-                        .map
-                        .entry(target_trait)
-                        .or_default()
-                        .entry(self_ty_fp)
-                        .or_default()
-                        .push(impl_id);
-                }
+    pub(crate) fn trait_impls_in_block_query(
+        db: &dyn HirDatabase,
+        block: BlockId,
+    ) -> Option<Arc<Self>> {
+        let _p = profile::span("trait_impls_in_block_query");
+        let mut impls = Self { map: FxHashMap::default() };
 
-                // To better support custom derives, collect impls in all unnamed const items.
-                // const _: () = { ... };
-                for konst in module_data.scope.unnamed_consts() {
-                    let body = db.body(konst.into());
-                    for (_, block_def_map) in body.blocks(db.upcast()) {
-                        collect_def_map(db, &block_def_map, impls);
-                    }
+        let block_def_map = db.block_def_map(block)?;
+        impls.collect_def_map(db, &block_def_map);
+
+        return Some(Arc::new(impls));
+    }
+
+    fn collect_def_map(&mut self, db: &dyn HirDatabase, def_map: &DefMap) {
+        for (_module_id, module_data) in def_map.modules() {
+            for impl_id in module_data.scope.impls() {
+                let target_trait = match db.impl_trait(impl_id) {
+                    Some(tr) => tr.skip_binders().hir_trait_id(),
+                    None => continue,
+                };
+                let self_ty = db.impl_self_ty(impl_id);
+                let self_ty_fp = TyFingerprint::for_trait_impl(self_ty.skip_binders());
+                self.map
+                    .entry(target_trait)
+                    .or_default()
+                    .entry(self_ty_fp)
+                    .or_default()
+                    .push(impl_id);
+            }
+
+            // To better support custom derives, collect impls in all unnamed const items.
+            // const _: () = { ... };
+            for konst in module_data.scope.unnamed_consts() {
+                let body = db.body(konst.into());
+                for (_, block_def_map) in body.blocks(db.upcast()) {
+                    self.collect_def_map(db, &block_def_map);
                 }
             }
         }