Auto merge of #14424 - Veykril:local-trait-impls, r=Veykril

fix: Properly handle local trait impls

Before we only handled trait impls that came from the block of either the trait or the target type, we now handle them correctly by tracking the block we are currently inferring from, then walking that up to collect all block trait impls.

Author: bors

crates/hir-def/src/db.rs

@@ -93,6 +93,8 @@ pub trait DefDatabase: InternDatabase + ExpandDatabase + Upcast<dyn ExpandDataba
     ///
     /// The `block_def_map` for block 0 would return `None`, while `block_def_map` of block 1 would
     /// return a `DefMap` containing `inner`.
+    // FIXME: This actually can't return None anymore as we no longer allocate block scopes for
+    // non item declaring blocks
     #[salsa::invoke(DefMap::block_def_map_query)]
     fn block_def_map(&self, block: BlockId) -> Option<Arc<DefMap>>;
 

crates/hir-ty/src/chalk_db.rs

@@ -1,8 +1,7 @@
 //! The implementation of `RustIrDatabase` for Chalk, which provides information
 //! about the code that Chalk needs.
-use std::sync::Arc;
+use std::{iter, sync::Arc};
 
-use cov_mark::hit;
 use tracing::debug;
 
 use chalk_ir::{cast::Cast, fold::shift::Shift, CanonicalVarKinds};
@@ -12,17 +11,16 @@ use base_db::CrateId;
 use hir_def::{
     expr::Movability,
     lang_item::{lang_attr, LangItem, LangItemTarget},
-    AssocItemId, GenericDefId, HasModule, ItemContainerId, Lookup, ModuleId, TypeAliasId,
+    AssocItemId, BlockId, GenericDefId, HasModule, ItemContainerId, Lookup, TypeAliasId,
 };
 use hir_expand::name::name;
 
 use crate::{
     db::HirDatabase,
     display::HirDisplay,
-    from_assoc_type_id, from_chalk_trait_id, from_foreign_def_id, make_binders,
-    make_single_type_binders,
+    from_assoc_type_id, from_chalk_trait_id, make_binders, make_single_type_binders,
     mapping::{from_chalk, ToChalk, TypeAliasAsValue},
-    method_resolution::{TraitImpls, TyFingerprint, ALL_FLOAT_FPS, ALL_INT_FPS},
+    method_resolution::{TyFingerprint, ALL_FLOAT_FPS, ALL_INT_FPS},
     to_assoc_type_id, to_chalk_trait_id,
     traits::ChalkContext,
     utils::generics,
@@ -108,53 +106,41 @@ impl<'a> chalk_solve::RustIrDatabase<Interner> for ChalkContext<'a> {
             _ => self_ty_fp.as_ref().map(std::slice::from_ref).unwrap_or(&[]),
         };
 
-        fn local_impls(db: &dyn HirDatabase, module: ModuleId) -> Option<Arc<TraitImpls>> {
-            let block = module.containing_block()?;
-            hit!(block_local_impls);
-            db.trait_impls_in_block(block)
-        }
-
         // 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 in_deps = self.db.trait_impls_in_deps(self.krate);
         let in_self = self.db.trait_impls_in_crate(self.krate);
-        let trait_module = trait_.module(self.db.upcast());
-        let type_module = match self_ty_fp {
-            Some(TyFingerprint::Adt(adt_id)) => Some(adt_id.module(self.db.upcast())),
-            Some(TyFingerprint::ForeignType(type_id)) => {
-                Some(from_foreign_def_id(type_id).module(self.db.upcast()))
-            }
-            Some(TyFingerprint::Dyn(trait_id)) => Some(trait_id.module(self.db.upcast())),
-            _ => None,
-        };
-        let impl_maps = [
-            Some(in_deps),
-            Some(in_self),
-            local_impls(self.db, trait_module),
-            type_module.and_then(|m| local_impls(self.db, m)),
-        ];
 
-        let id_to_chalk = |id: hir_def::ImplId| id.to_chalk(self.db);
+        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)
+                .and_then(|map| map.parent())
+                .and_then(|module| module.containing_block())
+        })
+        .filter_map(|block_id| self.db.trait_impls_in_block(block_id));
 
-        let result: Vec<_> = if fps.is_empty() {
-            debug!("Unrestricted search for {:?} impls...", trait_);
-            impl_maps
-                .iter()
-                .filter_map(|o| o.as_ref())
-                .flat_map(|impls| impls.for_trait(trait_).map(id_to_chalk))
-                .collect()
-        } else {
-            impl_maps
-                .iter()
-                .filter_map(|o| o.as_ref())
-                .flat_map(|impls| {
-                    fps.iter().flat_map(move |fp| {
-                        impls.for_trait_and_self_ty(trait_, *fp).map(id_to_chalk)
-                    })
-                })
-                .collect()
-        };
+        let id_to_chalk = |id: hir_def::ImplId| id.to_chalk(self.db);
+        let mut result = vec![];
+        match fps {
+            [] => {
+                debug!("Unrestricted search for {:?} impls...", trait_);
+                impl_maps.into_iter().chain(block_impls).for_each(|impls| {
+                    result.extend(impls.for_trait(trait_).map(id_to_chalk));
+                });
+            }
+            fps => {
+                impl_maps.into_iter().chain(block_impls).for_each(|impls| {
+                    result.extend(
+                        fps.iter().flat_map(|fp| {
+                            impls.for_trait_and_self_ty(trait_, *fp).map(id_to_chalk)
+                        }),
+                    );
+                });
+            }
+        }
 
         debug!("impls_for_trait returned {} impls", result.len());
         result
@@ -193,7 +179,7 @@ impl<'a> chalk_solve::RustIrDatabase<Interner> for ChalkContext<'a> {
         &self,
         environment: &chalk_ir::Environment<Interner>,
     ) -> chalk_ir::ProgramClauses<Interner> {
-        self.db.program_clauses_for_chalk_env(self.krate, environment.clone())
+        self.db.program_clauses_for_chalk_env(self.krate, self.block, environment.clone())
     }
 
     fn opaque_ty_data(&self, id: chalk_ir::OpaqueTyId<Interner>) -> Arc<OpaqueTyDatum> {
@@ -451,9 +437,10 @@ impl<'a> chalk_ir::UnificationDatabase<Interner> for &'a dyn HirDatabase {
 pub(crate) fn program_clauses_for_chalk_env_query(
     db: &dyn HirDatabase,
     krate: CrateId,
+    block: Option<BlockId>,
     environment: chalk_ir::Environment<Interner>,
 ) -> chalk_ir::ProgramClauses<Interner> {
-    chalk_solve::program_clauses_for_env(&ChalkContext { db, krate }, &environment)
+    chalk_solve::program_clauses_for_env(&ChalkContext { db, krate, block }, &environment)
 }
 
 pub(crate) fn associated_ty_data_query(

crates/hir-ty/src/db.rs

@@ -129,7 +129,7 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
     fn trait_impls_in_crate(&self, krate: CrateId) -> Arc<TraitImpls>;
 
     #[salsa::invoke(TraitImpls::trait_impls_in_block_query)]
-    fn trait_impls_in_block(&self, krate: BlockId) -> Option<Arc<TraitImpls>>;
+    fn trait_impls_in_block(&self, block: BlockId) -> Option<Arc<TraitImpls>>;
 
     #[salsa::invoke(TraitImpls::trait_impls_in_deps_query)]
     fn trait_impls_in_deps(&self, krate: CrateId) -> Arc<TraitImpls>;
@@ -197,20 +197,23 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
     fn trait_solve(
         &self,
         krate: CrateId,
+        block: Option<BlockId>,
         goal: crate::Canonical<crate::InEnvironment<crate::Goal>>,
     ) -> Option<crate::Solution>;
 
     #[salsa::invoke(crate::traits::trait_solve_query)]
     fn trait_solve_query(
         &self,
         krate: CrateId,
+        block: Option<BlockId>,
         goal: crate::Canonical<crate::InEnvironment<crate::Goal>>,
     ) -> Option<crate::Solution>;
 
     #[salsa::invoke(chalk_db::program_clauses_for_chalk_env_query)]
     fn program_clauses_for_chalk_env(
         &self,
         krate: CrateId,
+        block: Option<BlockId>,
         env: chalk_ir::Environment<Interner>,
     ) -> chalk_ir::ProgramClauses<Interner>;
 }
@@ -232,10 +235,11 @@ fn infer_wait(db: &dyn HirDatabase, def: DefWithBodyId) -> Arc<InferenceResult>
 fn trait_solve_wait(
     db: &dyn HirDatabase,
     krate: CrateId,
+    block: Option<BlockId>,
     goal: crate::Canonical<crate::InEnvironment<crate::Goal>>,
 ) -> Option<crate::Solution> {
     let _p = profile::span("trait_solve::wait");
-    db.trait_solve_query(krate, goal)
+    db.trait_solve_query(krate, block, goal)
 }
 
 #[test]

crates/hir-ty/src/infer.rs

@@ -40,7 +40,7 @@ use crate::{
     db::HirDatabase, fold_tys, fold_tys_and_consts, infer::coerce::CoerceMany,
     lower::ImplTraitLoweringMode, static_lifetime, to_assoc_type_id, AliasEq, AliasTy, Const,
     DomainGoal, GenericArg, Goal, ImplTraitId, InEnvironment, Interner, ProjectionTy, RpitId,
-    Substitution, TraitEnvironment, TraitRef, Ty, TyBuilder, TyExt, TyKind,
+    Substitution, TraitRef, Ty, TyBuilder, TyExt, TyKind,
 };
 
 // This lint has a false positive here. See the link below for details.
@@ -442,7 +442,6 @@ pub(crate) struct InferenceContext<'a> {
     pub(crate) body: &'a Body,
     pub(crate) resolver: Resolver,
     table: unify::InferenceTable<'a>,
-    trait_env: Arc<TraitEnvironment>,
     /// The traits in scope, disregarding block modules. This is used for caching purposes.
     traits_in_scope: FxHashSet<TraitId>,
     pub(crate) result: InferenceResult,
@@ -516,8 +515,7 @@ impl<'a> InferenceContext<'a> {
         let trait_env = db.trait_environment_for_body(owner);
         InferenceContext {
             result: InferenceResult::default(),
-            table: unify::InferenceTable::new(db, trait_env.clone()),
-            trait_env,
+            table: unify::InferenceTable::new(db, trait_env),
             return_ty: TyKind::Error.intern(Interner), // set in collect_* calls
             resume_yield_tys: None,
             return_coercion: None,

crates/hir-ty/src/infer/coerce.rs

@@ -636,7 +636,7 @@ impl<'a> InferenceTable<'a> {
         // Need to find out in what cases this is necessary
         let solution = self
             .db
-            .trait_solve(krate, canonicalized.value.clone().cast(Interner))
+            .trait_solve(krate, self.trait_env.block, canonicalized.value.clone().cast(Interner))
             .ok_or(TypeError)?;
 
         match solution {

crates/hir-ty/src/infer/expr.rs

@@ -3,6 +3,7 @@
 use std::{
     iter::{repeat, repeat_with},
     mem,
+    sync::Arc,
 };
 
 use chalk_ir::{
@@ -15,7 +16,7 @@ use hir_def::{
     generics::TypeOrConstParamData,
     lang_item::LangItem,
     path::{GenericArg, GenericArgs},
-    ConstParamId, FieldId, ItemContainerId, Lookup,
+    BlockId, ConstParamId, FieldId, ItemContainerId, Lookup,
 };
 use hir_expand::name::{name, Name};
 use stdx::always;
@@ -147,19 +148,19 @@ impl<'a> InferenceContext<'a> {
                 self.infer_top_pat(pat, &input_ty);
                 self.result.standard_types.bool_.clone()
             }
-            Expr::Block { statements, tail, label, id: _ } => {
-                self.infer_block(tgt_expr, statements, *tail, *label, expected)
+            Expr::Block { statements, tail, label, id } => {
+                self.infer_block(tgt_expr, *id, statements, *tail, *label, expected)
             }
-            Expr::Unsafe { id: _, statements, tail } => {
-                self.infer_block(tgt_expr, statements, *tail, None, expected)
+            Expr::Unsafe { id, statements, tail } => {
+                self.infer_block(tgt_expr, *id, statements, *tail, None, expected)
             }
-            Expr::Const { id: _, statements, tail } => {
+            Expr::Const { id, statements, tail } => {
                 self.with_breakable_ctx(BreakableKind::Border, None, None, |this| {
-                    this.infer_block(tgt_expr, statements, *tail, None, expected)
+                    this.infer_block(tgt_expr, *id, statements, *tail, None, expected)
                 })
                 .1
             }
-            Expr::Async { id: _, statements, tail } => {
+            Expr::Async { id, statements, tail } => {
                 let ret_ty = self.table.new_type_var();
                 let prev_diverges = mem::replace(&mut self.diverges, Diverges::Maybe);
                 let prev_ret_ty = mem::replace(&mut self.return_ty, ret_ty.clone());
@@ -170,6 +171,7 @@ impl<'a> InferenceContext<'a> {
                     self.with_breakable_ctx(BreakableKind::Border, None, None, |this| {
                         this.infer_block(
                             tgt_expr,
+                            *id,
                             statements,
                             *tail,
                             None,
@@ -394,7 +396,7 @@ impl<'a> InferenceContext<'a> {
                                 }
                             }
                             let trait_ = fn_x
-                                .get_id(self.db, self.trait_env.krate)
+                                .get_id(self.db, self.table.trait_env.krate)
                                 .expect("We just used it");
                             let trait_data = self.db.trait_data(trait_);
                             if let Some(func) = trait_data.method_by_name(&fn_x.method_name()) {
@@ -787,7 +789,7 @@ impl<'a> InferenceContext<'a> {
                     let canonicalized = self.canonicalize(base_ty.clone());
                     let receiver_adjustments = method_resolution::resolve_indexing_op(
                         self.db,
-                        self.trait_env.clone(),
+                        self.table.trait_env.clone(),
                         canonicalized.value,
                         index_trait,
                     );
@@ -1205,13 +1207,19 @@ impl<'a> InferenceContext<'a> {
     fn infer_block(
         &mut self,
         expr: ExprId,
+        block_id: Option<BlockId>,
         statements: &[Statement],
         tail: Option<ExprId>,
         label: Option<LabelId>,
         expected: &Expectation,
     ) -> Ty {
         let coerce_ty = expected.coercion_target_type(&mut self.table);
         let g = self.resolver.update_to_inner_scope(self.db.upcast(), self.owner, expr);
+        let prev_env = block_id.map(|block_id| {
+            let prev_env = self.table.trait_env.clone();
+            Arc::make_mut(&mut self.table.trait_env).block = Some(block_id);
+            prev_env
+        });
 
         let (break_ty, ty) =
             self.with_breakable_ctx(BreakableKind::Block, Some(coerce_ty.clone()), label, |this| {
@@ -1300,6 +1308,9 @@ impl<'a> InferenceContext<'a> {
                 }
             });
         self.resolver.reset_to_guard(g);
+        if let Some(prev_env) = prev_env {
+            self.table.trait_env = prev_env;
+        }
 
         break_ty.unwrap_or(ty)
     }
@@ -1398,7 +1409,7 @@ impl<'a> InferenceContext<'a> {
                     method_resolution::lookup_method(
                         self.db,
                         &canonicalized_receiver.value,
-                        self.trait_env.clone(),
+                        self.table.trait_env.clone(),
                         self.get_traits_in_scope().as_ref().left_or_else(|&it| it),
                         VisibleFromModule::Filter(self.resolver.module()),
                         name,
@@ -1431,7 +1442,7 @@ impl<'a> InferenceContext<'a> {
         let resolved = method_resolution::lookup_method(
             self.db,
             &canonicalized_receiver.value,
-            self.trait_env.clone(),
+            self.table.trait_env.clone(),
             self.get_traits_in_scope().as_ref().left_or_else(|&it| it),
             VisibleFromModule::Filter(self.resolver.module()),
             method_name,

crates/hir-ty/src/infer/unify.rs

@@ -462,7 +462,8 @@ impl<'a> InferenceTable<'a> {
     pub(crate) fn try_obligation(&mut self, goal: Goal) -> Option<Solution> {
         let in_env = InEnvironment::new(&self.trait_env.env, goal);
         let canonicalized = self.canonicalize(in_env);
-        let solution = self.db.trait_solve(self.trait_env.krate, canonicalized.value);
+        let solution =
+            self.db.trait_solve(self.trait_env.krate, self.trait_env.block, canonicalized.value);
         solution
     }
 
@@ -597,7 +598,11 @@ impl<'a> InferenceTable<'a> {
         &mut self,
         canonicalized: &Canonicalized<InEnvironment<Goal>>,
     ) -> bool {
-        let solution = self.db.trait_solve(self.trait_env.krate, canonicalized.value.clone());
+        let solution = self.db.trait_solve(
+            self.trait_env.krate,
+            self.trait_env.block,
+            canonicalized.value.clone(),
+        );
 
         match solution {
             Some(Solution::Unique(canonical_subst)) => {
@@ -684,7 +689,11 @@ impl<'a> InferenceTable<'a> {
             environment: trait_env.clone(),
         };
         let canonical = self.canonicalize(obligation.clone());
-        if self.db.trait_solve(krate, canonical.value.cast(Interner)).is_some() {
+        if self
+            .db
+            .trait_solve(krate, self.trait_env.block, canonical.value.cast(Interner))
+            .is_some()
+        {
             self.register_obligation(obligation.goal);
             let return_ty = self.normalize_projection_ty(projection);
             for fn_x in [FnTrait::Fn, FnTrait::FnMut, FnTrait::FnOnce] {
@@ -695,7 +704,11 @@ impl<'a> InferenceTable<'a> {
                     environment: trait_env.clone(),
                 };
                 let canonical = self.canonicalize(obligation.clone());
-                if self.db.trait_solve(krate, canonical.value.cast(Interner)).is_some() {
+                if self
+                    .db
+                    .trait_solve(krate, self.trait_env.block, canonical.value.cast(Interner))
+                    .is_some()
+                {
                     return Some((fn_x, arg_tys, return_ty));
                 }
             }

crates/hir-ty/src/lib.rs

@@ -1,6 +1,5 @@
 //! The type system. We currently use this to infer types for completion, hover
 //! information and various assists.
-
 #![warn(rust_2018_idioms, unused_lifetimes, semicolon_in_expressions_from_macros)]
 
 #[allow(unused)]