Allocate traits in scope upfront when type checking instead of recollecting them everytime

crates/hir-def/src/resolver.rs

@@ -449,6 +449,15 @@ impl Resolver {
         traits
     }
 
+    pub fn traits_in_scope_from_block_scopes(&self) -> impl Iterator<Item = TraitId> + '_ {
+        self.scopes()
+            .filter_map(|scope| match scope {
+                Scope::BlockScope(m) => Some(m.def_map[m.module_id].scope.traits()),
+                _ => None,
+            })
+            .flatten()
+    }
+
     pub fn module(&self) -> ModuleId {
         let (def_map, local_id) = self.item_scope();
         def_map.module_id(local_id)

crates/hir-ty/src/infer.rs

@@ -33,7 +33,7 @@ use hir_def::{
 };
 use hir_expand::name::{name, Name};
 use la_arena::ArenaMap;
-use rustc_hash::FxHashMap;
+use rustc_hash::{FxHashMap, FxHashSet};
 use stdx::always;
 
 use crate::{
@@ -423,6 +423,8 @@ pub(crate) struct InferenceContext<'a> {
     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,
     /// The return type of the function being inferred, the closure or async block if we're
     /// currently within one.
@@ -505,6 +507,7 @@ impl<'a> InferenceContext<'a> {
             db,
             owner,
             body,
+            traits_in_scope: resolver.traits_in_scope(db.upcast()),
             resolver,
             diverges: Diverges::Maybe,
             breakables: Vec::new(),
@@ -1056,6 +1059,15 @@ impl<'a> InferenceContext<'a> {
         let struct_ = self.resolve_lang_item(LangItem::VaList)?.as_struct()?;
         Some(struct_.into())
     }
+
+    fn get_traits_in_scope(&self) -> Either<FxHashSet<TraitId>, &FxHashSet<TraitId>> {
+        let mut b_traits = self.resolver.traits_in_scope_from_block_scopes().peekable();
+        if b_traits.peek().is_some() {
+            Either::Left(self.traits_in_scope.iter().copied().chain(b_traits).collect())
+        } else {
+            Either::Right(&self.traits_in_scope)
+        }
+    }
 }
 
 /// When inferring an expression, we propagate downward whatever type hint we

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

@@ -1349,14 +1349,14 @@ impl<'a> InferenceContext<'a> {
             None => {
                 // no field found,
                 let method_with_same_name_exists = {
-                    let canonicalized_receiver = self.canonicalize(receiver_ty.clone());
+                    self.get_traits_in_scope();
-                    let traits_in_scope = self.resolver.traits_in_scope(self.db.upcast());
 
+                    let canonicalized_receiver = self.canonicalize(receiver_ty.clone());
                     method_resolution::lookup_method(
                         self.db,
                         &canonicalized_receiver.value,
                         self.trait_env.clone(),
-                        &traits_in_scope,
+                        self.get_traits_in_scope().as_ref().left_or_else(|&it| it),
                         VisibleFromModule::Filter(self.resolver.module()),
                         name,
                     )
@@ -1385,13 +1385,11 @@ impl<'a> InferenceContext<'a> {
         let receiver_ty = self.infer_expr(receiver, &Expectation::none());
         let canonicalized_receiver = self.canonicalize(receiver_ty.clone());
 
-        let traits_in_scope = self.resolver.traits_in_scope(self.db.upcast());
-
         let resolved = method_resolution::lookup_method(
             self.db,
             &canonicalized_receiver.value,
             self.trait_env.clone(),
-            &traits_in_scope,
+            self.get_traits_in_scope().as_ref().left_or_else(|&it| it),
             VisibleFromModule::Filter(self.resolver.module()),
             method_name,
         );

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

@@ -205,6 +205,7 @@ impl<'a> InferenceContext<'a> {
         Some((def, Some(trait_ref.substitution)))
     }
 
+    // FIXME: Change sig to -> Option<(ValueNs, Substitution)>, subs aren't optional from here anymore
     fn resolve_ty_assoc_item(
         &mut self,
         ty: Ty,
@@ -220,25 +221,35 @@ impl<'a> InferenceContext<'a> {
         }
 
         let canonical_ty = self.canonicalize(ty.clone());
-        let traits_in_scope = self.resolver.traits_in_scope(self.db.upcast());
 
         let mut not_visible = None;
         let res = method_resolution::iterate_method_candidates(
             &canonical_ty.value,
             self.db,
             self.table.trait_env.clone(),
-            &traits_in_scope,
+            self.get_traits_in_scope().as_ref().left_or_else(|&it| it),
             VisibleFromModule::Filter(self.resolver.module()),
             Some(name),
             method_resolution::LookupMode::Path,
             |_ty, item, visible| {
+                if visible {
+                    Some((item, true))
+                } else {
+                    if not_visible.is_none() {
+                        not_visible = Some((item, false));
+                    }
+                    None
+                }
+            },
+        );
+        let res = res.or(not_visible);
+        let (item, visible) = res?;
+
         let (def, container) = match item {
             AssocItemId::FunctionId(f) => {
                 (ValueNs::FunctionId(f), f.lookup(self.db.upcast()).container)
             }
-                    AssocItemId::ConstId(c) => {
+            AssocItemId::ConstId(c) => (ValueNs::ConstId(c), c.lookup(self.db.upcast()).container),
-                        (ValueNs::ConstId(c), c.lookup(self.db.upcast()).container)
-                    }
             AssocItemId::TypeAliasId(_) => unreachable!(),
         };
         let substs = match container {
@@ -246,8 +257,7 @@ impl<'a> InferenceContext<'a> {
                 let impl_substs = TyBuilder::subst_for_def(self.db, impl_id, None)
                     .fill_with_inference_vars(&mut self.table)
                     .build();
-                        let impl_self_ty =
+                let impl_self_ty = self.db.impl_self_ty(impl_id).substitute(Interner, &impl_substs);
-                            self.db.impl_self_ty(impl_id).substitute(Interner, &impl_substs);
                 self.unify(&impl_self_ty, &ty);
                 impl_substs
             }
@@ -266,24 +276,11 @@ impl<'a> InferenceContext<'a> {
             }
         };
 
-                if visible {
-                    Some((def, item, Some(substs), true))
-                } else {
-                    if not_visible.is_none() {
-                        not_visible = Some((def, item, Some(substs), false));
-                    }
-                    None
-                }
-            },
-        );
-        let res = res.or(not_visible);
-        if let Some((_, item, Some(ref substs), visible)) = res {
         self.write_assoc_resolution(id, item, substs.clone());
         if !visible {
-                self.push_diagnostic(InferenceDiagnostic::PrivateAssocItem { id, item })
+            self.push_diagnostic(InferenceDiagnostic::PrivateAssocItem { id, item });
         }
-        }
+        Some((def, Some(substs)))
-        res.map(|(def, _, substs, _)| (def, substs))
     }
 
     fn resolve_enum_variant_on_ty(