3 files changed, 104 insertions, 59 deletions

diff --git a/crates/hir_ty/src/infer.rs b/crates/hir_ty/src/infer.rs
index 531159e54..1d8184ef5 100644
--- a/crates/hir_ty/src/infer.rs
+++ b/crates/hir_ty/src/infer.rs
@@ -18,7 +18,7 @@ use std::mem;
 use std::ops::Index;
 use std::sync::Arc;
 
-use chalk_ir::{cast::Cast, Mutability};
+use chalk_ir::{cast::Cast, DebruijnIndex, Mutability};
 use hir_def::{
     body::Body,
     data::{ConstData, FunctionData, StaticData},
@@ -41,8 +41,9 @@ use super::{
     TypeWalk,
 };
 use crate::{
-    db::HirDatabase, infer::diagnostics::InferenceDiagnostic, lower::ImplTraitLoweringMode,
-    to_assoc_type_id, AliasEq, AliasTy, Canonical, Interner, TyBuilder, TyExt, TyKind,
+    db::HirDatabase, fold_tys, infer::diagnostics::InferenceDiagnostic,
+    lower::ImplTraitLoweringMode, to_assoc_type_id, AliasEq, AliasTy, Canonical, Interner,
+    TyBuilder, TyExt, TyKind,
 };
 
 // This lint has a false positive here. See the link below for details.
@@ -323,7 +324,7 @@ impl<'a> InferenceContext<'a> {
     }
 
     fn insert_type_vars(&mut self, ty: Ty) -> Ty {
-        ty.fold(&mut |ty| self.insert_type_vars_shallow(ty))
+        fold_tys(ty, |ty, _| self.insert_type_vars_shallow(ty), DebruijnIndex::INNERMOST)
     }
 
     fn resolve_obligations_as_possible(&mut self) {
@@ -434,12 +435,16 @@ impl<'a> InferenceContext<'a> {
     /// to do it as well.
     fn normalize_associated_types_in(&mut self, ty: Ty) -> Ty {
         let ty = self.resolve_ty_as_possible(ty);
-        ty.fold(&mut |ty| match ty.kind(&Interner) {
-            TyKind::Alias(AliasTy::Projection(proj_ty)) => {
-                self.normalize_projection_ty(proj_ty.clone())
-            }
-            _ => 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 {
diff --git a/crates/hir_ty/src/infer/unify.rs b/crates/hir_ty/src/infer/unify.rs
index c3d9e81c3..ffa324b6d 100644
--- a/crates/hir_ty/src/infer/unify.rs
+++ b/crates/hir_ty/src/infer/unify.rs
@@ -3,15 +3,15 @@
 use std::borrow::Cow;
 
 use chalk_ir::{
-    interner::HasInterner, FloatTy, IntTy, TyVariableKind, UniverseIndex, VariableKind,
+    fold::Fold, interner::HasInterner, FloatTy, IntTy, TyVariableKind, UniverseIndex, VariableKind,
 };
 use ena::unify::{InPlaceUnificationTable, NoError, UnifyKey, UnifyValue};
 
 use super::{DomainGoal, InferenceContext};
 use crate::{
-    AliasEq, AliasTy, BoundVar, Canonical, CanonicalVarKinds, DebruijnIndex, FnPointer, FnSubst,
-    InEnvironment, InferenceVar, Interner, Scalar, Substitution, Ty, TyExt, TyKind, TypeWalk,
-    WhereClause,
+    fold_tys, AliasEq, AliasTy, BoundVar, Canonical, CanonicalVarKinds, DebruijnIndex, FnPointer,
+    FnSubst, InEnvironment, InferenceVar, Interner, Scalar, Substitution, Ty, TyExt, TyKind,
+    TypeWalk, WhereClause,
 };
 
 impl<'a> InferenceContext<'a> {
@@ -53,9 +53,14 @@ impl<'a, 'b> Canonicalizer<'a, 'b> {
         })
     }
 
-    fn do_canonicalize<T: TypeWalk>(&mut self, t: T, binders: DebruijnIndex) -> T {
-        t.fold_binders(
-            &mut |ty, binders| match ty.kind(&Interner) {
+    fn do_canonicalize<T: Fold<Interner, Result = T> + HasInterner<Interner = Interner>>(
+        &mut self,
+        t: T,
+        binders: DebruijnIndex,
+    ) -> T {
+        fold_tys(
+            t,
+            |ty, binders| match ty.kind(&Interner) {
                 &TyKind::InferenceVar(var, kind) => {
                     let inner = from_inference_var(var);
                     if self.var_stack.contains(&inner) {
@@ -485,55 +490,63 @@ impl InferenceTable {
     /// be resolved as far as possible, i.e. contain no type variables with
     /// known type.
     fn resolve_ty_as_possible_inner(&mut self, tv_stack: &mut Vec<TypeVarId>, ty: Ty) -> Ty {
-        ty.fold(&mut |ty| match ty.kind(&Interner) {
-            &TyKind::InferenceVar(tv, kind) => {
-                let inner = from_inference_var(tv);
-                if tv_stack.contains(&inner) {
-                    cov_mark::hit!(type_var_cycles_resolve_as_possible);
-                    // recursive type
-                    return self.type_variable_table.fallback_value(tv, kind);
-                }
-                if let Some(known_ty) =
-                    self.var_unification_table.inlined_probe_value(inner).known()
-                {
-                    // known_ty may contain other variables that are known by now
-                    tv_stack.push(inner);
-                    let result = self.resolve_ty_as_possible_inner(tv_stack, known_ty.clone());
-                    tv_stack.pop();
-                    result
-                } else {
-                    ty
+        fold_tys(
+            ty,
+            |ty, _| match ty.kind(&Interner) {
+                &TyKind::InferenceVar(tv, kind) => {
+                    let inner = from_inference_var(tv);
+                    if tv_stack.contains(&inner) {
+                        cov_mark::hit!(type_var_cycles_resolve_as_possible);
+                        // recursive type
+                        return self.type_variable_table.fallback_value(tv, kind);
+                    }
+                    if let Some(known_ty) =
+                        self.var_unification_table.inlined_probe_value(inner).known()
+                    {
+                        // known_ty may contain other variables that are known by now
+                        tv_stack.push(inner);
+                        let result = self.resolve_ty_as_possible_inner(tv_stack, known_ty.clone());
+                        tv_stack.pop();
+                        result
+                    } else {
+                        ty
+                    }
                 }
-            }
-            _ => ty,
-        })
+                _ => ty,
+            },
+            DebruijnIndex::INNERMOST,
+        )
     }
 
     /// Resolves the type completely; type variables without known type are
     /// replaced by TyKind::Unknown.
     fn resolve_ty_completely_inner(&mut self, tv_stack: &mut Vec<TypeVarId>, ty: Ty) -> Ty {
-        ty.fold(&mut |ty| match ty.kind(&Interner) {
-            &TyKind::InferenceVar(tv, kind) => {
-                let inner = from_inference_var(tv);
-                if tv_stack.contains(&inner) {
-                    cov_mark::hit!(type_var_cycles_resolve_completely);
-                    // recursive type
-                    return self.type_variable_table.fallback_value(tv, kind);
-                }
-                if let Some(known_ty) =
-                    self.var_unification_table.inlined_probe_value(inner).known()
-                {
-                    // known_ty may contain other variables that are known by now
-                    tv_stack.push(inner);
-                    let result = self.resolve_ty_completely_inner(tv_stack, known_ty.clone());
-                    tv_stack.pop();
-                    result
-                } else {
-                    self.type_variable_table.fallback_value(tv, kind)
+        fold_tys(
+            ty,
+            |ty, _| match ty.kind(&Interner) {
+                &TyKind::InferenceVar(tv, kind) => {
+                    let inner = from_inference_var(tv);
+                    if tv_stack.contains(&inner) {
+                        cov_mark::hit!(type_var_cycles_resolve_completely);
+                        // recursive type
+                        return self.type_variable_table.fallback_value(tv, kind);
+                    }
+                    if let Some(known_ty) =
+                        self.var_unification_table.inlined_probe_value(inner).known()
+                    {
+                        // known_ty may contain other variables that are known by now
+                        tv_stack.push(inner);
+                        let result = self.resolve_ty_completely_inner(tv_stack, known_ty.clone());
+                        tv_stack.pop();
+                        result
+                    } else {
+                        self.type_variable_table.fallback_value(tv, kind)
+                    }
                 }
-            }
-            _ => ty,
-        })
+                _ => ty,
+            },
+            DebruijnIndex::INNERMOST,
+        )
     }
 }
 
diff --git a/crates/hir_ty/src/lib.rs b/crates/hir_ty/src/lib.rs
index 60680ec86..f7af8b8aa 100644
--- a/crates/hir_ty/src/lib.rs
+++ b/crates/hir_ty/src/lib.rs
@@ -328,3 +328,30 @@ pub(crate) fn fold_free_vars<T: HasInterner<Interner = Interner> + Fold<Interner
     }
     t.fold_with(&mut FreeVarFolder(f), DebruijnIndex::INNERMOST).expect("fold failed unexpectedly")
 }
+
+pub(crate) fn fold_tys<T: HasInterner<Interner = Interner> + Fold<Interner>>(
+    t: T,
+    f: impl FnMut(Ty, DebruijnIndex) -> Ty,
+    binders: DebruijnIndex,
+) -> T::Result {
+    use chalk_ir::{
+        fold::{Folder, SuperFold},
+        Fallible,
+    };
+    struct TyFolder<F>(F);
+    impl<'i, F: FnMut(Ty, DebruijnIndex) -> Ty + 'i> Folder<'i, Interner> for TyFolder<F> {
+        fn as_dyn(&mut self) -> &mut dyn Folder<'i, Interner> {
+            self
+        }
+
+        fn interner(&self) -> &'i Interner {
+            &Interner
+        }
+
+        fn fold_ty(&mut self, ty: Ty, outer_binder: DebruijnIndex) -> Fallible<Ty> {
+            let ty = ty.super_fold_with(self.as_dyn(), outer_binder)?;
+            Ok(self.0(ty, outer_binder))
+        }
+    }
+    t.fold_with(&mut TyFolder(f), binders).expect("fold failed unexpectedly")
+}