1 files changed, 40 insertions, 16 deletions

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));
         }
     }
 }