Make canonicalization API a bit nicer

2 files changed, 56 insertions, 31 deletions

diff --git a/crates/ra_hir/src/ty/infer.rs b/crates/ra_hir/src/ty/infer.rs
index 41ae569f7..cc74c6322 100644
--- a/crates/ra_hir/src/ty/infer.rs
+++ b/crates/ra_hir/src/ty/infer.rs
@@ -324,23 +324,25 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
     fn resolve_obligations_as_possible(&mut self) {
         let obligations = mem::replace(&mut self.obligations, Vec::new());
         for obligation in obligations {
-            let mut canonicalizer = self.canonicalizer();
-            let solution = match &obligation {
+            let (solution, canonicalized) = match &obligation {
                 Obligation::Trait(tr) => {
-                    let canonical = canonicalizer.canonicalize_trait_ref(tr.clone());
-                    super::traits::implements(
-                        canonicalizer.ctx.db,
-                        canonicalizer.ctx.resolver.krate().unwrap(),
-                        canonical,
+                    let canonicalized = self.canonicalizer().canonicalize_trait_ref(tr.clone());
+                    (
+                        super::traits::implements(
+                            self.db,
+                            self.resolver.krate().unwrap(),
+                            canonicalized.value.clone(),
+                        ),
+                        canonicalized,
                     )
                 }
             };
             match solution {
                 Some(Solution::Unique(substs)) => {
-                    canonicalizer.apply_solution(substs.0);
+                    canonicalized.apply_solution(self, substs.0);
                 }
                 Some(Solution::Ambig(Guidance::Definite(substs))) => {
-                    canonicalizer.apply_solution(substs.0);
+                    canonicalized.apply_solution(self, substs.0);
                     self.obligations.push(obligation);
                 }
                 Some(_) => {
@@ -877,17 +879,16 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
         generic_args: Option<&GenericArgs>,
     ) -> Ty {
         let receiver_ty = self.infer_expr(receiver, &Expectation::none());
-        let mut canonicalizer = self.canonicalizer();
-        let canonical_receiver = canonicalizer.canonicalize_ty(receiver_ty.clone());
+        let canonicalized_receiver = self.canonicalizer().canonicalize_ty(receiver_ty.clone());
         let resolved = method_resolution::lookup_method(
-            &canonical_receiver,
-            canonicalizer.ctx.db,
+            &canonicalized_receiver.value,
+            self.db,
             method_name,
-            &canonicalizer.ctx.resolver,
+            &self.resolver,
         );
         let (derefed_receiver_ty, method_ty, def_generics) = match resolved {
             Some((ty, func)) => {
-                let ty = canonicalizer.decanonicalize_ty(ty);
+                let ty = canonicalized_receiver.decanonicalize_ty(ty);
                 self.write_method_resolution(tgt_expr, func);
                 (
                     ty,
diff --git a/crates/ra_hir/src/ty/infer/unify.rs b/crates/ra_hir/src/ty/infer/unify.rs
index 820a64789..c9c8959c4 100644
--- a/crates/ra_hir/src/ty/infer/unify.rs
+++ b/crates/ra_hir/src/ty/infer/unify.rs
@@ -13,14 +13,17 @@ impl<'a, D: HirDatabase> InferenceContext<'a, D> {
     }
 }
 
-// TODO improve the interface of this
-
 pub(super) struct Canonicalizer<'a, 'b, D: HirDatabase>
 where
     'a: 'b,
 {
-    pub ctx: &'b mut InferenceContext<'a, D>,
-    pub free_vars: Vec<InferTy>,
+    ctx: &'b mut InferenceContext<'a, D>,
+    free_vars: Vec<InferTy>,
+}
+
+pub(super) struct Canonicalized<T> {
+    pub value: Canonical<T>,
+    free_vars: Vec<InferTy>,
 }
 
 impl<'a, 'b, D: HirDatabase> Canonicalizer<'a, 'b, D>
@@ -35,13 +38,13 @@ where
         })
     }
 
-    pub fn canonicalize_ty(&mut self, ty: Ty) -> Canonical<Ty> {
-        let value = ty.fold(&mut |ty| match ty {
+    fn do_canonicalize_ty(&mut self, ty: Ty) -> Ty {
+        ty.fold(&mut |ty| match ty {
             Ty::Infer(tv) => {
                 let inner = tv.to_inner();
                 // TODO prevent infinite loops? => keep var stack
                 if let Some(known_ty) = self.ctx.var_unification_table.probe_value(inner).known() {
-                    self.canonicalize_ty(known_ty.clone()).value
+                    self.do_canonicalize_ty(known_ty.clone())
                 } else {
                     let free_var = InferTy::TypeVar(self.ctx.var_unification_table.find(inner));
                     let position = self.add(free_var);
@@ -49,20 +52,37 @@ where
                 }
             }
             _ => ty,
-        });
-        Canonical { value, num_vars: self.free_vars.len() }
+        })
     }
 
-    pub fn canonicalize_trait_ref(&mut self, trait_ref: TraitRef) -> Canonical<TraitRef> {
+    fn do_canonicalize_trait_ref(&mut self, trait_ref: TraitRef) -> TraitRef {
         let substs = trait_ref
             .substs
             .iter()
-            .map(|ty| self.canonicalize_ty(ty.clone()).value)
+            .map(|ty| self.do_canonicalize_ty(ty.clone()))
             .collect::<Vec<_>>();
-        let value = TraitRef { trait_: trait_ref.trait_, substs: substs.into() };
-        Canonical { value, num_vars: self.free_vars.len() }
+        TraitRef { trait_: trait_ref.trait_, substs: substs.into() }
     }
 
+    fn into_canonicalized<T>(self, result: T) -> Canonicalized<T> {
+        Canonicalized {
+            value: Canonical { value: result, num_vars: self.free_vars.len() },
+            free_vars: self.free_vars,
+        }
+    }
+
+    pub fn canonicalize_ty(mut self, ty: Ty) -> Canonicalized<Ty> {
+        let result = self.do_canonicalize_ty(ty);
+        self.into_canonicalized(result)
+    }
+
+    pub fn canonicalize_trait_ref(mut self, trait_ref: TraitRef) -> Canonicalized<TraitRef> {
+        let result = self.do_canonicalize_trait_ref(trait_ref);
+        self.into_canonicalized(result)
+    }
+}
+
+impl<T> Canonicalized<T> {
     pub fn decanonicalize_ty(&self, ty: Ty) -> Ty {
         ty.fold(&mut |ty| match ty {
             Ty::Bound(idx) => {
@@ -76,13 +96,17 @@ where
         })
     }
 
-    pub fn apply_solution(&mut self, solution: Canonical<Vec<Ty>>) {
+    pub fn apply_solution(
+        &self,
+        ctx: &mut InferenceContext<'_, impl HirDatabase>,
+        solution: Canonical<Vec<Ty>>,
+    ) {
         // the solution may contain new variables, which we need to convert to new inference vars
         let new_vars =
-            (0..solution.num_vars).map(|_| self.ctx.new_type_var()).collect::<Vec<_>>().into();
+            (0..solution.num_vars).map(|_| ctx.new_type_var()).collect::<Vec<_>>().into();
         for (i, ty) in solution.value.into_iter().enumerate() {
             let var = self.free_vars[i].clone();
-            self.ctx.unify(&Ty::Infer(var), &ty.subst_bound_vars(&new_vars));
+            ctx.unify(&Ty::Infer(var), &ty.subst_bound_vars(&new_vars));
         }
     }
 }