Rework obligation handling

We can't do the easy hack that we did before anymore, where we kept
track of whether any inference variables changed since the last time we
rechecked obligations. Instead, we store the obligations in
canonicalized form; that way we can easily check the inference variables
to see whether they have changed since the goal was canonicalized.

11 files changed, 240 insertions, 143 deletions

diff --git a/crates/hir/src/lib.rs b/crates/hir/src/lib.rs
index 1429384cb..52d72c3c5 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -1712,15 +1712,17 @@ impl Type {
         resolver: &Resolver,
         ty: Ty,
     ) -> Type {
-        let environment =
-            resolver.generic_def().map_or_else(Default::default, |d| db.trait_environment(d));
+        let environment = resolver
+            .generic_def()
+            .map_or_else(|| Arc::new(TraitEnvironment::empty(krate)), |d| db.trait_environment(d));
         Type { krate, env: environment, ty }
     }
 
     fn new(db: &dyn HirDatabase, krate: CrateId, lexical_env: impl HasResolver, ty: Ty) -> Type {
         let resolver = lexical_env.resolver(db.upcast());
-        let environment =
-            resolver.generic_def().map_or_else(Default::default, |d| db.trait_environment(d));
+        let environment = resolver
+            .generic_def()
+            .map_or_else(|| Arc::new(TraitEnvironment::empty(krate)), |d| db.trait_environment(d));
         Type { krate, env: environment, ty }
     }
 
diff --git a/crates/hir_ty/src/db.rs b/crates/hir_ty/src/db.rs
index f773aa621..be5b9110e 100644
--- a/crates/hir_ty/src/db.rs
+++ b/crates/hir_ty/src/db.rs
@@ -134,14 +134,14 @@ pub trait HirDatabase: DefDatabase + Upcast<dyn DefDatabase> {
     fn trait_solve(
         &self,
         krate: CrateId,
-        goal: crate::Canonical<crate::InEnvironment<crate::DomainGoal>>,
+        goal: crate::Canonical<crate::InEnvironment<crate::Goal>>,
     ) -> Option<crate::Solution>;
 
     #[salsa::invoke(crate::traits::trait_solve_query)]
     fn trait_solve_query(
         &self,
         krate: CrateId,
-        goal: crate::Canonical<crate::InEnvironment<crate::DomainGoal>>,
+        goal: crate::Canonical<crate::InEnvironment<crate::Goal>>,
     ) -> Option<crate::Solution>;
 
     #[salsa::invoke(chalk_db::program_clauses_for_chalk_env_query)]
@@ -168,7 +168,7 @@ fn infer_wait(db: &dyn HirDatabase, def: DefWithBodyId) -> Arc<InferenceResult>
 fn trait_solve_wait(
     db: &dyn HirDatabase,
     krate: CrateId,
-    goal: crate::Canonical<crate::InEnvironment<crate::DomainGoal>>,
+    goal: crate::Canonical<crate::InEnvironment<crate::Goal>>,
 ) -> Option<crate::Solution> {
     let _p = profile::span("trait_solve::wait");
     db.trait_solve_query(krate, goal)
diff --git a/crates/hir_ty/src/infer.rs b/crates/hir_ty/src/infer.rs
index 7898740be..97e7c5f8c 100644
--- a/crates/hir_ty/src/infer.rs
+++ b/crates/hir_ty/src/infer.rs
@@ -14,7 +14,7 @@
 //! the `ena` crate, which is extracted from rustc.
 
 use std::borrow::Cow;
-use std::mem;
+
 use std::ops::Index;
 use std::sync::Arc;
 
@@ -27,8 +27,8 @@ use hir_def::{
     path::{path, Path},
     resolver::{HasResolver, Resolver, TypeNs},
     type_ref::TypeRef,
-    AdtId, AssocItemId, DefWithBodyId, EnumVariantId, FieldId, FunctionId, Lookup, TraitId,
-    TypeAliasId, VariantId,
+    AdtId, AssocItemId, DefWithBodyId, EnumVariantId, FieldId, FunctionId, HasModule, Lookup,
+    TraitId, TypeAliasId, VariantId,
 };
 use hir_expand::{diagnostics::DiagnosticSink, name::name};
 use la_arena::ArenaMap;
@@ -36,13 +36,11 @@ use rustc_hash::FxHashMap;
 use stdx::impl_from;
 use syntax::SmolStr;
 
-use super::{
-    DomainGoal, Guidance, InEnvironment, ProjectionTy, Solution, TraitEnvironment, TraitRef, Ty,
-};
+use super::{DomainGoal, InEnvironment, ProjectionTy, TraitEnvironment, TraitRef, Ty};
 use crate::{
     db::HirDatabase, fold_tys, infer::diagnostics::InferenceDiagnostic,
-    lower::ImplTraitLoweringMode, to_assoc_type_id, AliasEq, AliasTy, Canonical, Interner,
-    TyBuilder, TyExt, TyKind,
+    lower::ImplTraitLoweringMode, to_assoc_type_id, AliasEq, AliasTy, Interner, TyBuilder, TyExt,
+    TyKind,
 };
 
 // This lint has a false positive here. See the link below for details.
@@ -227,8 +225,6 @@ struct InferenceContext<'a> {
     resolver: Resolver,
     table: unify::InferenceTable<'a>,
     trait_env: Arc<TraitEnvironment>,
-    obligations: Vec<DomainGoal>,
-    last_obligations_check: Option<u32>,
     result: InferenceResult,
     /// The return type of the function being inferred, or the closure if we're
     /// currently within one.
@@ -260,15 +256,15 @@ fn find_breakable<'c>(
 
 impl<'a> InferenceContext<'a> {
     fn new(db: &'a dyn HirDatabase, owner: DefWithBodyId, resolver: Resolver) -> Self {
-        let trait_env =
-            owner.as_generic_def_id().map_or_else(Default::default, |d| db.trait_environment(d));
+        let krate = owner.module(db.upcast()).krate();
+        let trait_env = owner
+            .as_generic_def_id()
+            .map_or_else(|| Arc::new(TraitEnvironment::empty(krate)), |d| db.trait_environment(d));
         InferenceContext {
             result: InferenceResult::default(),
             table: unify::InferenceTable::new(db, trait_env.clone()),
-            obligations: Vec::default(),
-            last_obligations_check: None,
-            return_ty: TyKind::Error.intern(&Interner), // set in collect_fn_signature
             trait_env,
+            return_ty: TyKind::Error.intern(&Interner), // set in collect_fn_signature
             db,
             owner,
             body: db.body(owner),
@@ -284,6 +280,7 @@ impl<'a> InferenceContext<'a> {
 
     fn resolve_all(mut self) -> InferenceResult {
         // FIXME resolve obligations as well (use Guidance if necessary)
+        self.table.resolve_obligations_as_possible();
 
         // make sure diverging type variables are marked as such
         self.table.propagate_diverging_flag();
@@ -357,44 +354,11 @@ impl<'a> InferenceContext<'a> {
     }
 
     fn resolve_obligations_as_possible(&mut self) {
-        let _span = profile::span("resolve_obligations_as_possible");
-
-        let obligations = mem::replace(&mut self.obligations, Vec::new());
-        for obligation in obligations {
-            let in_env = InEnvironment::new(&self.trait_env.env, obligation.clone());
-            let canonicalized = self.canonicalize(in_env);
-            let solution =
-                self.db.trait_solve(self.resolver.krate().unwrap(), canonicalized.value.clone());
-
-            match solution {
-                Some(Solution::Unique(canonical_subst)) => {
-                    canonicalized.apply_solution(
-                        self,
-                        Canonical {
-                            binders: canonical_subst.binders,
-                            // FIXME: handle constraints
-                            value: canonical_subst.value.subst,
-                        },
-                    );
-                }
-                Some(Solution::Ambig(Guidance::Definite(substs))) => {
-                    canonicalized.apply_solution(self, substs);
-                    self.obligations.push(obligation);
-                }
-                Some(_) => {
-                    // FIXME use this when trying to resolve everything at the end
-                    self.obligations.push(obligation);
-                }
-                None => {
-                    // FIXME obligation cannot be fulfilled => diagnostic
-                }
-            };
-        }
+        self.table.resolve_obligations_as_possible();
     }
 
     fn push_obligation(&mut self, o: DomainGoal) {
-        self.obligations.push(o);
-        self.last_obligations_check = None;
+        self.table.register_obligation(o.cast(&Interner));
     }
 
     fn unify(&mut self, ty1: &Ty, ty2: &Ty) -> bool {
@@ -467,25 +431,7 @@ impl<'a> InferenceContext<'a> {
     /// call). `make_ty` handles this already, but e.g. for field types we need
     /// to do it as well.
     fn normalize_associated_types_in(&mut self, ty: Ty) -> Ty {
-        let ty = self.resolve_ty_as_possible(ty);
-        fold_tys(
-            ty,
-            |ty, _| match ty.kind(&Interner) {
-                TyKind::Alias(AliasTy::Projection(proj_ty)) => {
-                    self.normalize_projection_ty(proj_ty.clone())
-                }
-                _ => ty,
-            },
-            DebruijnIndex::INNERMOST,
-        )
-    }
-
-    fn normalize_projection_ty(&mut self, proj_ty: ProjectionTy) -> Ty {
-        let var = self.table.new_type_var();
-        let alias_eq = AliasEq { alias: AliasTy::Projection(proj_ty), ty: var.clone() };
-        let obligation = alias_eq.cast(&Interner);
-        self.push_obligation(obligation);
-        var
+        self.table.normalize_associated_types_in(ty)
     }
 
     fn resolve_variant(&mut self, path: Option<&Path>) -> (Ty, Option<VariantId>) {
diff --git a/crates/hir_ty/src/infer/coerce.rs b/crates/hir_ty/src/infer/coerce.rs
index 4d80b4a08..911343cb9 100644
--- a/crates/hir_ty/src/infer/coerce.rs
+++ b/crates/hir_ty/src/infer/coerce.rs
@@ -402,12 +402,15 @@ impl<'a> InferenceContext<'a> {
         // solve `CoerceUnsized` and `Unsize` goals at this point and leaves the
         // rest for later. Also, there's some logic about sized type variables.
         // Need to find out in what cases this is necessary
-        let solution = self.db.trait_solve(krate, canonicalized.value.clone()).ok_or(TypeError)?;
+        let solution = self
+            .db
+            .trait_solve(krate, canonicalized.value.clone().cast(&Interner))
+            .ok_or(TypeError)?;
 
         match solution {
             Solution::Unique(v) => {
                 canonicalized.apply_solution(
-                    self,
+                    &mut self.table,
                     Canonical {
                         binders: v.binders,
                         // FIXME handle constraints
diff --git a/crates/hir_ty/src/infer/expr.rs b/crates/hir_ty/src/infer/expr.rs
index f439169ea..6eaccd9b4 100644
--- a/crates/hir_ty/src/infer/expr.rs
+++ b/crates/hir_ty/src/infer/expr.rs
@@ -99,9 +99,9 @@ impl<'a> InferenceContext<'a> {
             environment: trait_env,
         };
         let canonical = self.canonicalize(obligation.clone());
-        if self.db.trait_solve(krate, canonical.value).is_some() {
+        if self.db.trait_solve(krate, canonical.value.cast(&Interner)).is_some() {
             self.push_obligation(obligation.goal);
-            let return_ty = self.normalize_projection_ty(projection);
+            let return_ty = self.table.normalize_projection_ty(projection);
             Some((arg_tys, return_ty))
         } else {
             None
@@ -306,7 +306,7 @@ impl<'a> InferenceContext<'a> {
                     self.resolver.krate(),
                     InEnvironment {
                         goal: canonicalized.value.clone(),
-                        environment: self.trait_env.env.clone(),
+                        environment: self.table.trait_env.env.clone(),
                     },
                 );
                 let (param_tys, ret_ty): (Vec<Ty>, Ty) = derefs
diff --git a/crates/hir_ty/src/infer/path.rs b/crates/hir_ty/src/infer/path.rs
index bc64b612b..fd366e121 100644
--- a/crates/hir_ty/src/infer/path.rs
+++ b/crates/hir_ty/src/infer/path.rs
@@ -225,7 +225,7 @@ impl<'a> InferenceContext<'a> {
         method_resolution::iterate_method_candidates(
             &canonical_ty.value,
             self.db,
-            self.trait_env.clone(),
+            self.table.trait_env.clone(),
             krate,
             &traits_in_scope,
             None,
diff --git a/crates/hir_ty/src/infer/unify.rs b/crates/hir_ty/src/infer/unify.rs
index 93cd54f0d..75e04e8b5 100644
--- a/crates/hir_ty/src/infer/unify.rs
+++ b/crates/hir_ty/src/infer/unify.rs
@@ -1,6 +1,6 @@
 //! Unification and canonicalization logic.
 
-use std::{borrow::Cow, fmt, sync::Arc};
+use std::{borrow::Cow, fmt, mem, sync::Arc};
 
 use chalk_ir::{
     cast::Cast, fold::Fold, interner::HasInterner, zip::Zip, FloatTy, IntTy, TyVariableKind,
@@ -11,8 +11,9 @@ use ena::unify::UnifyKey;
 
 use super::{InferOk, InferResult, InferenceContext, TypeError};
 use crate::{
-    db::HirDatabase, fold_tys, static_lifetime, BoundVar, Canonical, DebruijnIndex, GenericArg,
-    InferenceVar, Interner, Scalar, Substitution, TraitEnvironment, Ty, TyKind, VariableKind,
+    db::HirDatabase, fold_tys, static_lifetime, AliasEq, AliasTy, BoundVar, Canonical,
+    DebruijnIndex, GenericArg, Goal, Guidance, InEnvironment, InferenceVar, Interner, ProjectionTy,
+    Scalar, Solution, Substitution, TraitEnvironment, Ty, TyKind, VariableKind,
 };
 
 impl<'a> InferenceContext<'a> {
@@ -23,17 +24,11 @@ impl<'a> InferenceContext<'a> {
     where
         T::Result: HasInterner<Interner = Interner>,
     {
-        let result = self.table.var_unification_table.canonicalize(&Interner, t);
-        let free_vars = result
-            .free_vars
-            .into_iter()
-            .map(|free_var| free_var.to_generic_arg(&Interner))
-            .collect();
-        Canonicalized { value: result.quantified, free_vars }
+        self.table.canonicalize(t)
     }
 }
 
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub(super) struct Canonicalized<T>
 where
     T: HasInterner<Interner = Interner>,
@@ -49,22 +44,16 @@ impl<T: HasInterner<Interner = Interner>> Canonicalized<T> {
 
     pub(super) fn apply_solution(
         &self,
-        ctx: &mut InferenceContext<'_>,
+        ctx: &mut InferenceTable,
         solution: Canonical<Substitution>,
     ) {
         // the solution may contain new variables, which we need to convert to new inference vars
         let new_vars = Substitution::from_iter(
             &Interner,
             solution.binders.iter(&Interner).map(|k| match k.kind {
-                VariableKind::Ty(TyVariableKind::General) => {
-                    ctx.table.new_type_var().cast(&Interner)
-                }
-                VariableKind::Ty(TyVariableKind::Integer) => {
-                    ctx.table.new_integer_var().cast(&Interner)
-                }
-                VariableKind::Ty(TyVariableKind::Float) => {
-                    ctx.table.new_float_var().cast(&Interner)
-                }
+                VariableKind::Ty(TyVariableKind::General) => ctx.new_type_var().cast(&Interner),
+                VariableKind::Ty(TyVariableKind::Integer) => ctx.new_integer_var().cast(&Interner),
+                VariableKind::Ty(TyVariableKind::Float) => ctx.new_float_var().cast(&Interner),
                 // Chalk can sometimes return new lifetime variables. We just use the static lifetime everywhere
                 VariableKind::Lifetime => static_lifetime().cast(&Interner),
                 _ => panic!("const variable in solution"),
@@ -76,9 +65,9 @@ impl<T: HasInterner<Interner = Interner>> Canonicalized<T> {
                 // eagerly replace projections in the type; we may be getting types
                 // e.g. from where clauses where this hasn't happened yet
                 let ty = ctx.normalize_associated_types_in(new_vars.apply(ty.clone(), &Interner));
-                ctx.table.unify(var.assert_ty_ref(&Interner), &ty);
+                ctx.unify(var.assert_ty_ref(&Interner), &ty);
             } else {
-                let _ = ctx.table.unify_inner(&var, &new_vars.apply(v.clone(), &Interner));
+                let _ = ctx.unify_inner(&var, &new_vars.apply(v.clone(), &Interner));
             }
         }
     }
@@ -167,10 +156,11 @@ type ChalkInferenceTable = chalk_solve::infer::InferenceTable<Interner>;
 
 #[derive(Clone)]
 pub(crate) struct InferenceTable<'a> {
-    db: &'a dyn HirDatabase,
-    trait_env: Arc<TraitEnvironment>,
+    pub db: &'a dyn HirDatabase,
+    pub trait_env: Arc<TraitEnvironment>,
     pub(super) var_unification_table: ChalkInferenceTable,
     pub(super) type_variable_table: TypeVariableTable,
+    pending_obligations: Vec<Canonicalized<InEnvironment<Goal>>>,
 }
 
 impl<'a> InferenceTable<'a> {
@@ -180,6 +170,7 @@ impl<'a> InferenceTable<'a> {
             trait_env,
             var_unification_table: ChalkInferenceTable::new(),
             type_variable_table: TypeVariableTable { inner: Vec::new() },
+            pending_obligations: Vec::new(),
         }
     }
 
@@ -202,6 +193,50 @@ impl<'a> InferenceTable<'a> {
         }
     }
 
+    pub(super) fn canonicalize<T: Fold<Interner> + HasInterner<Interner = Interner>>(
+        &mut self,
+        t: T,
+    ) -> Canonicalized<T::Result>
+    where
+        T::Result: HasInterner<Interner = Interner>,
+    {
+        let result = self.var_unification_table.canonicalize(&Interner, t);
+        let free_vars = result
+            .free_vars
+            .into_iter()
+            .map(|free_var| free_var.to_generic_arg(&Interner))
+            .collect();
+        Canonicalized { value: result.quantified, free_vars }
+    }
+
+    /// Recurses through the given type, normalizing associated types mentioned
+    /// in it by replacing them by type variables and registering obligations to
+    /// resolve later. This should be done once for every type we get from some
+    /// type annotation (e.g. from a let type annotation, field type or function
+    /// call). `make_ty` handles this already, but e.g. for field types we need
+    /// to do it as well.
+    pub(super) fn normalize_associated_types_in(&mut self, ty: Ty) -> Ty {
+        let ty = self.resolve_ty_as_possible(ty);
+        fold_tys(
+            ty,
+            |ty, _| match ty.kind(&Interner) {
+                TyKind::Alias(AliasTy::Projection(proj_ty)) => {
+                    self.normalize_projection_ty(proj_ty.clone())
+                }
+                _ => ty,
+            },
+            DebruijnIndex::INNERMOST,
+        )
+    }
+
+    pub(super) fn normalize_projection_ty(&mut self, proj_ty: ProjectionTy) -> Ty {
+        let var = self.new_type_var();
+        let alias_eq = AliasEq { alias: AliasTy::Projection(proj_ty), ty: var.clone() };
+        let obligation = alias_eq.cast(&Interner);
+        self.register_obligation(obligation);
+        var
+    }
+
     fn new_var(&mut self, kind: TyVariableKind, diverging: bool) -> Ty {
         let var = self.var_unification_table.new_variable(UniverseIndex::ROOT);
         // Chalk might have created some type variables for its own purposes that we don't know about...
@@ -341,6 +376,94 @@ impl<'a> InferenceTable<'a> {
             DebruijnIndex::INNERMOST,
         )
     }
+
+    pub fn register_obligation(&mut self, goal: Goal) {
+        let in_env = InEnvironment::new(&self.trait_env.env, goal);
+        self.register_obligation_in_env(in_env)
+    }
+
+    fn register_obligation_in_env(&mut self, goal: InEnvironment<Goal>) {
+        let canonicalized = self.canonicalize(goal);
+        if !self.try_resolve_obligation(&canonicalized) {
+            self.pending_obligations.push(canonicalized);
+        }
+    }
+
+    pub fn resolve_obligations_as_possible(&mut self) {
+        let _span = profile::span("resolve_obligations_as_possible");
+        let mut changed = true;
+        let mut obligations = Vec::new();
+        while changed {
+            changed = false;
+            mem::swap(&mut self.pending_obligations, &mut obligations);
+            for canonicalized in obligations.drain(..) {
+                if !self.check_changed(&canonicalized) {
+                    self.pending_obligations.push(canonicalized);
+                    continue;
+                }
+                changed = true;
+                let uncanonical = chalk_ir::Substitute::apply(
+                    &canonicalized.free_vars,
+                    canonicalized.value.value,
+                    &Interner,
+                );
+                self.register_obligation_in_env(uncanonical);
+            }
+        }
+    }
+
+    /// This checks whether any of the free variables in the `canonicalized`
+    /// have changed (either been unified with another variable, or with a
+    /// value). If this is not the case, we don't need to try to solve the goal
+    /// again -- it'll give the same result as last time.
+    fn check_changed(&mut self, canonicalized: &Canonicalized<InEnvironment<Goal>>) -> bool {
+        canonicalized.free_vars.iter().any(|var| {
+            let iv = match var.data(&Interner) {
+                chalk_ir::GenericArgData::Ty(ty) => ty.inference_var(&Interner),
+                chalk_ir::GenericArgData::Lifetime(lt) => lt.inference_var(&Interner),
+                chalk_ir::GenericArgData::Const(c) => c.inference_var(&Interner),
+            }
+            .expect("free var is not inference var");
+            if self.var_unification_table.probe_var(iv).is_some() {
+                return true;
+            }
+            let root = self.var_unification_table.inference_var_root(iv);
+            iv != root
+        })
+    }
+
+    fn try_resolve_obligation(
+        &mut self,
+        canonicalized: &Canonicalized<InEnvironment<Goal>>,
+    ) -> bool {
+        let solution = self.db.trait_solve(self.trait_env.krate, canonicalized.value.clone());
+
+        match solution {
+            Some(Solution::Unique(canonical_subst)) => {
+                canonicalized.apply_solution(
+                    self,
+                    Canonical {
+                        binders: canonical_subst.binders,
+                        // FIXME: handle constraints
+                        value: canonical_subst.value.subst,
+                    },
+                );
+                true
+            }
+            Some(Solution::Ambig(Guidance::Definite(substs))) => {
+                canonicalized.apply_solution(self, substs);
+                false
+            }
+            Some(_) => {
+                // FIXME use this when trying to resolve everything at the end
+                false
+            }
+            None => {
+                // FIXME obligation cannot be fulfilled => diagnostic
+                true
+            }
+        }
+    }
 }
 
 impl<'a> fmt::Debug for InferenceTable<'a> {
diff --git a/crates/hir_ty/src/lib.rs b/crates/hir_ty/src/lib.rs
index 56f60c46b..72093d75a 100644
--- a/crates/hir_ty/src/lib.rs
+++ b/crates/hir_ty/src/lib.rs
@@ -45,7 +45,7 @@ use hir_def::{
 };
 use stdx::always;
 
-use crate::{db::HirDatabase, display::HirDisplay, utils::generics};
+use crate::{db::HirDatabase, utils::generics};
 
 pub use autoderef::autoderef;
 pub use builder::TyBuilder;
@@ -114,6 +114,7 @@ pub type FnSig = chalk_ir::FnSig<Interner>;
 
 pub type InEnvironment<T> = chalk_ir::InEnvironment<T>;
 pub type DomainGoal = chalk_ir::DomainGoal<Interner>;
+pub type Goal = chalk_ir::Goal<Interner>;
 pub type AliasEq = chalk_ir::AliasEq<Interner>;
 pub type Solution = chalk_solve::Solution<Interner>;
 pub type ConstrainedSubst = chalk_ir::ConstrainedSubst<Interner>;
diff --git a/crates/hir_ty/src/lower.rs b/crates/hir_ty/src/lower.rs
index bd8bb6028..8a375b973 100644
--- a/crates/hir_ty/src/lower.rs
+++ b/crates/hir_ty/src/lower.rs
@@ -1035,9 +1035,11 @@ pub(crate) fn trait_environment_query(
         clauses.push(program_clause.into_from_env_clause(&Interner));
     }
 
+    let krate = def.module(db.upcast()).krate();
+
     let env = chalk_ir::Environment::new(&Interner).add_clauses(&Interner, clauses);
 
-    Arc::new(TraitEnvironment { traits_from_clauses: traits_in_scope, env })
+    Arc::new(TraitEnvironment { krate, traits_from_clauses: traits_in_scope, env })
 }
 
 /// Resolve the where clause(s) of an item with generics.
diff --git a/crates/hir_ty/src/method_resolution.rs b/crates/hir_ty/src/method_resolution.rs
index 37e1f89d8..08e385a42 100644
--- a/crates/hir_ty/src/method_resolution.rs
+++ b/crates/hir_ty/src/method_resolution.rs
@@ -577,6 +577,7 @@ fn iterate_method_candidates_by_receiver(
         if iterate_inherent_methods(
             self_ty,
             db,
+            env.clone(),
             name,
             Some(receiver_ty),
             krate,
@@ -613,8 +614,16 @@ fn iterate_method_candidates_for_self_ty(
     name: Option<&Name>,
     mut callback: &mut dyn FnMut(&Ty, AssocItemId) -> bool,
 ) -> bool {
-    if iterate_inherent_methods(self_ty, db, name, None, krate, visible_from_module, &mut callback)
-    {
+    if iterate_inherent_methods(
+        self_ty,
+        db,
+        env.clone(),
+        name,
+        None,
+        krate,
+        visible_from_module,
+        &mut callback,
+    ) {
         return true;
     }
     iterate_trait_method_candidates(self_ty, db, env, krate, traits_in_scope, name, None, callback)
@@ -653,12 +662,12 @@ fn iterate_trait_method_candidates(
         for (_name, 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.
-            if !is_valid_candidate(db, name, receiver_ty, *item, self_ty, None) {
+            if !is_valid_candidate(db, env.clone(), name, receiver_ty, *item, self_ty, None) {
                 continue;
             }
             if !known_implemented {
                 let goal = generic_implements_goal(db, env.clone(), t, self_ty.clone());
-                if db.trait_solve(krate, goal).is_none() {
+                if db.trait_solve(krate, goal.cast(&Interner)).is_none() {
                     continue 'traits;
                 }
             }
@@ -675,6 +684,7 @@ fn iterate_trait_method_candidates(
 fn iterate_inherent_methods(
     self_ty: &Canonical<Ty>,
     db: &dyn HirDatabase,
+    env: Arc<TraitEnvironment>,
     name: Option<&Name>,
     receiver_ty: Option<&Canonical<Ty>>,
     krate: CrateId,
@@ -690,14 +700,24 @@ fn iterate_inherent_methods(
 
         for &impl_def in impls.for_self_ty(&self_ty.value) {
             for &item in db.impl_data(impl_def).items.iter() {
-                if !is_valid_candidate(db, name, receiver_ty, item, self_ty, visible_from_module) {
+                if !is_valid_candidate(
+                    db,
+                    env.clone(),
+                    name,
+                    receiver_ty,
+                    item,
+                    self_ty,
+                    visible_from_module,
+                ) {
                     continue;
                 }
                 // we have to check whether the self type unifies with the type
                 // that the impl is for. If we have a receiver type, this
                 // already happens in `is_valid_candidate` above; if not, we
                 // check it here
-                if receiver_ty.is_none() && inherent_impl_substs(db, impl_def, self_ty).is_none() {
+                if receiver_ty.is_none()
+                    && inherent_impl_substs(db, env.clone(), impl_def, self_ty).is_none()
+                {
                     cov_mark::hit!(impl_self_type_match_without_receiver);
                     continue;
                 }
@@ -722,7 +742,7 @@ pub fn resolve_indexing_op(
     let deref_chain = autoderef_method_receiver(db, krate, ty);
     for ty in deref_chain {
         let goal = generic_implements_goal(db, env.clone(), index_trait, ty.clone());
-        if db.trait_solve(krate, goal).is_some() {
+        if db.trait_solve(krate, goal.cast(&Interner)).is_some() {
             return Some(ty);
         }
     }
@@ -731,6 +751,7 @@ pub fn resolve_indexing_op(
 
 fn is_valid_candidate(
     db: &dyn HirDatabase,
+    env: Arc<TraitEnvironment>,
     name: Option<&Name>,
     receiver_ty: Option<&Canonical<Ty>>,
     item: AssocItemId,
@@ -749,7 +770,7 @@ fn is_valid_candidate(
                 if !data.has_self_param() {
                     return false;
                 }
-                let transformed_receiver_ty = match transform_receiver_ty(db, m, self_ty) {
+                let transformed_receiver_ty = match transform_receiver_ty(db, env, m, self_ty) {
                     Some(ty) => ty,
                     None => return false,
                 };
@@ -776,6 +797,7 @@ fn is_valid_candidate(
 
 pub(crate) fn inherent_impl_substs(
     db: &dyn HirDatabase,
+    env: Arc<TraitEnvironment>,
     impl_id: ImplId,
     self_ty: &Canonical<Ty>,
 ) -> Option<Substitution> {
@@ -798,8 +820,7 @@ pub(crate) fn inherent_impl_substs(
         binders: CanonicalVarKinds::from_iter(&Interner, kinds),
         value: (self_ty_with_vars, self_ty.value.clone()),
     };
-    let trait_env = Arc::new(TraitEnvironment::default()); // FIXME
-    let substs = super::infer::unify(db, trait_env, &tys)?;
+    let substs = super::infer::unify(db, env, &tys)?;
     // We only want the substs for the vars we added, not the ones from self_ty.
     // Also, if any of the vars we added are still in there, we replace them by
     // Unknown. I think this can only really happen if self_ty contained
@@ -824,6 +845,7 @@ fn fallback_bound_vars(s: Substitution, num_vars_to_keep: usize) -> Substitution
 
 fn transform_receiver_ty(
     db: &dyn HirDatabase,
+    env: Arc<TraitEnvironment>,
     function_id: FunctionId,
     self_ty: &Canonical<Ty>,
 ) -> Option<Ty> {
@@ -833,7 +855,7 @@ fn transform_receiver_ty(
             .fill_with_unknown()
             .build(),
         AssocContainerId::ImplId(impl_id) => {
-            let impl_substs = inherent_impl_substs(db, impl_id, &self_ty)?;
+            let impl_substs = inherent_impl_substs(db, env, impl_id, &self_ty)?;
             TyBuilder::subst_for_def(db, function_id)
                 .use_parent_substs(&impl_substs)
                 .fill_with_unknown()
@@ -853,7 +875,7 @@ pub fn implements_trait(
     trait_: TraitId,
 ) -> bool {
     let goal = generic_implements_goal(db, env, trait_, ty.clone());
-    let solution = db.trait_solve(krate, goal);
+    let solution = db.trait_solve(krate, goal.cast(&Interner));
 
     solution.is_some()
 }
@@ -866,7 +888,7 @@ pub fn implements_trait_unique(
     trait_: TraitId,
 ) -> bool {
     let goal = generic_implements_goal(db, env, trait_, ty.clone());
-    let solution = db.trait_solve(krate, goal);
+    let solution = db.trait_solve(krate, goal.cast(&Interner));
 
     matches!(solution, Some(crate::Solution::Unique(_)))
 }
diff --git a/crates/hir_ty/src/traits.rs b/crates/hir_ty/src/traits.rs
index 9936d0803..294cb531c 100644
--- a/crates/hir_ty/src/traits.rs
+++ b/crates/hir_ty/src/traits.rs
@@ -2,7 +2,7 @@
 
 use std::env::var;
 
-use chalk_ir::cast::Cast;
+use chalk_ir::GoalData;
 use chalk_solve::{logging_db::LoggingRustIrDatabase, Solver};
 
 use base_db::CrateId;
@@ -10,7 +10,7 @@ use hir_def::{lang_item::LangItemTarget, TraitId};
 use stdx::panic_context;
 
 use crate::{
-    db::HirDatabase, AliasEq, AliasTy, Canonical, DomainGoal, Guidance, HirDisplay, InEnvironment,
+    db::HirDatabase, AliasEq, AliasTy, Canonical, DomainGoal, Goal, Guidance, InEnvironment,
     Interner, Solution, TraitRefExt, Ty, TyKind, WhereClause,
 };
 
@@ -38,6 +38,7 @@ fn create_chalk_solver() -> chalk_recursive::RecursiveSolver<Interner> {
 /// we assume that `T: Default`.
 #[derive(Debug, Clone, PartialEq, Eq, Hash)]
 pub struct TraitEnvironment {
+    pub krate: CrateId,
     // When we're using Chalk's Ty we can make this a BTreeMap since it's Ord,
     // but for now it's too annoying...
     pub(crate) traits_from_clauses: Vec<(Ty, TraitId)>,
@@ -45,6 +46,14 @@ pub struct TraitEnvironment {
 }
 
 impl TraitEnvironment {
+    pub fn empty(krate: CrateId) -> Self {
+        TraitEnvironment {
+            krate,
+            traits_from_clauses: Vec::new(),
+            env: chalk_ir::Environment::new(&Interner),
+        }
+    }
+
     pub(crate) fn traits_in_scope_from_clauses<'a>(
         &'a self,
         ty: &'a Ty,
@@ -59,34 +68,25 @@ impl TraitEnvironment {
     }
 }
 
-impl Default for TraitEnvironment {
-    fn default() -> Self {
-        TraitEnvironment {
-            traits_from_clauses: Vec::new(),
-            env: chalk_ir::Environment::new(&Interner),
-        }
-    }
-}
-
 /// Solve a trait goal using Chalk.
 pub(crate) fn trait_solve_query(
     db: &dyn HirDatabase,
     krate: CrateId,
-    goal: Canonical<InEnvironment<DomainGoal>>,
+    goal: Canonical<InEnvironment<Goal>>,
 ) -> Option<Solution> {
-    let _p = profile::span("trait_solve_query").detail(|| match &goal.value.goal {
-        DomainGoal::Holds(WhereClause::Implemented(it)) => {
+    let _p = profile::span("trait_solve_query").detail(|| match &goal.value.goal.data(&Interner) {
+        GoalData::DomainGoal(DomainGoal::Holds(WhereClause::Implemented(it))) => {
             db.trait_data(it.hir_trait_id()).name.to_string()
         }
-        DomainGoal::Holds(WhereClause::AliasEq(_)) => "alias_eq".to_string(),
+        GoalData::DomainGoal(DomainGoal::Holds(WhereClause::AliasEq(_))) => "alias_eq".to_string(),
         _ => "??".to_string(),
     });
-    log::info!("trait_solve_query({})", goal.value.goal.display(db));
+    log::info!("trait_solve_query({:?})", goal.value.goal);
 
-    if let DomainGoal::Holds(WhereClause::AliasEq(AliasEq {
+    if let GoalData::DomainGoal(DomainGoal::Holds(WhereClause::AliasEq(AliasEq {
         alias: AliasTy::Projection(projection_ty),
         ..
-    })) = &goal.value.goal
+    }))) = &goal.value.goal.data(&Interner)
     {
         if let TyKind::BoundVar(_) = projection_ty.self_type_parameter(&Interner).kind(&Interner) {
             // Hack: don't ask Chalk to normalize with an unknown self type, it'll say that's impossible
@@ -94,11 +94,9 @@ pub(crate) fn trait_solve_query(
         }
     }
 
-    let canonical = goal.cast(&Interner);
-
     // We currently don't deal with universes (I think / hope they're not yet
     // relevant for our use cases?)
-    let u_canonical = chalk_ir::UCanonical { canonical, universes: 1 };
+    let u_canonical = chalk_ir::UCanonical { canonical: goal, universes: 1 };
     solve(db, krate, &u_canonical)
 }