6 files changed, 76 insertions, 50 deletions

diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 571b89bc3..cda050a7d 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -54,8 +54,8 @@ use hir_ty::{
     method_resolution, to_assoc_type_id,
     traits::{FnTrait, Solution, SolutionVariables},
     AliasTy, BoundVar, CallableDefId, CallableSig, Canonical, DebruijnIndex, GenericPredicate,
-    InEnvironment, Interner, Obligation, ProjectionPredicate, ProjectionTy, Scalar, Substs,
-    TraitEnvironment, Ty, TyDefId, TyKind, TyVariableKind,
+    InEnvironment, Interner, Obligation, ProjectionPredicate, ProjectionTy, Scalar, Substs, Ty,
+    TyDefId, TyKind, TyVariableKind,
 };
 use rustc_hash::FxHashSet;
 use stdx::{format_to, impl_from};
@@ -817,7 +817,7 @@ impl Function {
         let resolver = self.id.resolver(db.upcast());
         let krate = self.id.lookup(db.upcast()).container.module(db.upcast()).krate();
         let ctx = hir_ty::TyLoweringContext::new(db, &resolver);
-        let environment = TraitEnvironment::lower(db, &resolver);
+        let environment = db.trait_environment(self.id.into());
         db.function_data(self.id)
             .params
             .iter()
@@ -1563,13 +1563,15 @@ impl Type {
         resolver: &Resolver,
         ty: Ty,
     ) -> Type {
-        let environment = TraitEnvironment::lower(db, &resolver);
+        let environment =
+            resolver.generic_def().map_or_else(Default::default, |d| db.trait_environment(d));
         Type { krate, ty: InEnvironment { value: ty, environment } }
     }
 
     fn new(db: &dyn HirDatabase, krate: CrateId, lexical_env: impl HasResolver, ty: Ty) -> Type {
         let resolver = lexical_env.resolver(db.upcast());
-        let environment = TraitEnvironment::lower(db, &resolver);
+        let environment =
+            resolver.generic_def().map_or_else(Default::default, |d| db.trait_environment(d));
         Type { krate, ty: InEnvironment { value: ty, environment } }
     }
 
diff --git a/crates/hir_def/src/lib.rs b/crates/hir_def/src/lib.rs
index 6d11c5be4..c6655c5fb 100644
--- a/crates/hir_def/src/lib.rs
+++ b/crates/hir_def/src/lib.rs
@@ -341,6 +341,16 @@ pub enum DefWithBodyId {
 
 impl_from!(FunctionId, ConstId, StaticId for DefWithBodyId);
 
+impl DefWithBodyId {
+    pub fn as_generic_def_id(self) -> Option<GenericDefId> {
+        match self {
+            DefWithBodyId::FunctionId(f) => Some(f.into()),
+            DefWithBodyId::StaticId(_) => None,
+            DefWithBodyId::ConstId(c) => Some(c.into()),
+        }
+    }
+}
+
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Hash)]
 pub enum AssocItemId {
     FunctionId(FunctionId),
diff --git a/crates/hir_ty/src/db.rs b/crates/hir_ty/src/db.rs
index 8a3cc0283..74a048672 100644
--- a/crates/hir_ty/src/db.rs
+++ b/crates/hir_ty/src/db.rs
@@ -65,6 +65,9 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
     #[salsa::invoke(crate::lower::generic_predicates_query)]
     fn generic_predicates(&self, def: GenericDefId) -> Arc<[Binders<GenericPredicate>]>;
 
+    #[salsa::invoke(crate::lower::trait_environment_query)]
+    fn trait_environment(&self, def: GenericDefId) -> Arc<crate::TraitEnvironment>;
+
     #[salsa::invoke(crate::lower::generic_defaults_query)]
     fn generic_defaults(&self, def: GenericDefId) -> Arc<[Binders<Ty>]>;
 
diff --git a/crates/hir_ty/src/infer.rs b/crates/hir_ty/src/infer.rs
index 4f7463422..bc52f447d 100644
--- a/crates/hir_ty/src/infer.rs
+++ b/crates/hir_ty/src/infer.rs
@@ -228,7 +228,9 @@ impl<'a> InferenceContext<'a> {
             table: unify::InferenceTable::new(),
             obligations: Vec::default(),
             return_ty: TyKind::Unknown.intern(&Interner), // set in collect_fn_signature
-            trait_env: TraitEnvironment::lower(db, &resolver),
+            trait_env: owner
+                .as_generic_def_id()
+                .map_or_else(Default::default, |d| db.trait_environment(d)),
             db,
             owner,
             body: db.body(owner),
diff --git a/crates/hir_ty/src/lower.rs b/crates/hir_ty/src/lower.rs
index e57d5970f..3b0706530 100644
--- a/crates/hir_ty/src/lower.rs
+++ b/crates/hir_ty/src/lower.rs
@@ -935,55 +935,55 @@ pub(crate) fn generic_predicates_for_param_recover(
     Arc::new([])
 }
 
-impl TraitEnvironment {
-    pub fn lower(db: &dyn HirDatabase, resolver: &Resolver) -> Arc<TraitEnvironment> {
-        let ctx = TyLoweringContext::new(db, &resolver)
-            .with_type_param_mode(TypeParamLoweringMode::Placeholder);
-        let mut traits_in_scope = Vec::new();
-        let mut clauses = Vec::new();
-        for pred in resolver.where_predicates_in_scope() {
-            for pred in GenericPredicate::from_where_predicate(&ctx, pred) {
-                if pred.is_error() {
-                    continue;
-                }
-                if let GenericPredicate::Implemented(tr) = &pred {
-                    traits_in_scope.push((tr.self_ty().clone(), tr.trait_));
-                }
-                let program_clause: chalk_ir::ProgramClause<Interner> =
-                    pred.clone().to_chalk(db).cast(&Interner);
-                clauses.push(program_clause.into_from_env_clause(&Interner));
+pub(crate) fn trait_environment_query(
+    db: &dyn HirDatabase,
+    def: GenericDefId,
+) -> Arc<TraitEnvironment> {
+    let resolver = def.resolver(db.upcast());
+    let ctx = TyLoweringContext::new(db, &resolver)
+        .with_type_param_mode(TypeParamLoweringMode::Placeholder);
+    let mut traits_in_scope = Vec::new();
+    let mut clauses = Vec::new();
+    for pred in resolver.where_predicates_in_scope() {
+        for pred in GenericPredicate::from_where_predicate(&ctx, pred) {
+            if pred.is_error() {
+                continue;
             }
-        }
-
-        if let Some(def) = resolver.generic_def() {
-            let container: Option<AssocContainerId> = match def {
-                // FIXME: is there a function for this?
-                GenericDefId::FunctionId(f) => Some(f.lookup(db.upcast()).container),
-                GenericDefId::AdtId(_) => None,
-                GenericDefId::TraitId(_) => None,
-                GenericDefId::TypeAliasId(t) => Some(t.lookup(db.upcast()).container),
-                GenericDefId::ImplId(_) => None,
-                GenericDefId::EnumVariantId(_) => None,
-                GenericDefId::ConstId(c) => Some(c.lookup(db.upcast()).container),
-            };
-            if let Some(AssocContainerId::TraitId(trait_id)) = container {
-                // add `Self: Trait<T1, T2, ...>` to the environment in trait
-                // function default implementations (and hypothetical code
-                // inside consts or type aliases)
-                cov_mark::hit!(trait_self_implements_self);
-                let substs = Substs::type_params(db, trait_id);
-                let trait_ref = TraitRef { trait_: trait_id, substs };
-                let pred = GenericPredicate::Implemented(trait_ref);
-                let program_clause: chalk_ir::ProgramClause<Interner> =
-                    pred.clone().to_chalk(db).cast(&Interner);
-                clauses.push(program_clause.into_from_env_clause(&Interner));
+            if let GenericPredicate::Implemented(tr) = &pred {
+                traits_in_scope.push((tr.self_ty().clone(), tr.trait_));
             }
+            let program_clause: chalk_ir::ProgramClause<Interner> =
+                pred.clone().to_chalk(db).cast(&Interner);
+            clauses.push(program_clause.into_from_env_clause(&Interner));
         }
+    }
 
-        let env = chalk_ir::Environment::new(&Interner).add_clauses(&Interner, clauses);
-
-        Arc::new(TraitEnvironment { traits_from_clauses: traits_in_scope, env })
+    let container: Option<AssocContainerId> = match def {
+        // FIXME: is there a function for this?
+        GenericDefId::FunctionId(f) => Some(f.lookup(db.upcast()).container),
+        GenericDefId::AdtId(_) => None,
+        GenericDefId::TraitId(_) => None,
+        GenericDefId::TypeAliasId(t) => Some(t.lookup(db.upcast()).container),
+        GenericDefId::ImplId(_) => None,
+        GenericDefId::EnumVariantId(_) => None,
+        GenericDefId::ConstId(c) => Some(c.lookup(db.upcast()).container),
+    };
+    if let Some(AssocContainerId::TraitId(trait_id)) = container {
+        // add `Self: Trait<T1, T2, ...>` to the environment in trait
+        // function default implementations (and hypothetical code
+        // inside consts or type aliases)
+        cov_mark::hit!(trait_self_implements_self);
+        let substs = Substs::type_params(db, trait_id);
+        let trait_ref = TraitRef { trait_: trait_id, substs };
+        let pred = GenericPredicate::Implemented(trait_ref);
+        let program_clause: chalk_ir::ProgramClause<Interner> =
+            pred.clone().to_chalk(db).cast(&Interner);
+        clauses.push(program_clause.into_from_env_clause(&Interner));
     }
+
+    let env = chalk_ir::Environment::new(&Interner).add_clauses(&Interner, clauses);
+
+    Arc::new(TraitEnvironment { traits_from_clauses: traits_in_scope, env })
 }
 
 /// Resolve the where clause(s) of an item with generics.
diff --git a/crates/hir_ty/src/traits.rs b/crates/hir_ty/src/traits.rs
index 500d1781c..edfafdff8 100644
--- a/crates/hir_ty/src/traits.rs
+++ b/crates/hir_ty/src/traits.rs
@@ -63,6 +63,15 @@ impl TraitEnvironment {
     }
 }
 
+impl Default for TraitEnvironment {
+    fn default() -> Self {
+        TraitEnvironment {
+            traits_from_clauses: Vec::new(),
+            env: chalk_ir::Environment::new(&Interner),
+        }
+    }
+}
+
 /// Something (usually a goal), along with an environment.
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct InEnvironment<T> {