resolve before canonicalization, ICE if unresolved

3 files changed, 115 insertions, 68 deletions

diff --git a/compiler/rustc_middle/src/traits/solve.rs b/compiler/rustc_middle/src/traits/solve.rs
index b21a00e4122..90139f925ed 100644
--- a/compiler/rustc_middle/src/traits/solve.rs
+++ b/compiler/rustc_middle/src/traits/solve.rs
@@ -147,7 +147,7 @@ impl<'tcx> std::ops::Deref for ExternalConstraints<'tcx> {
 }
 
 /// Additional constraints returned on success.
-#[derive(Debug, PartialEq, Eq, Clone, Hash, HashStable, Default)]
+#[derive(Debug, PartialEq, Eq, Clone, Hash, HashStable, Default, TypeVisitable, TypeFoldable)]
 pub struct ExternalConstraintsData<'tcx> {
     // FIXME: implement this.
     pub region_constraints: QueryRegionConstraints<'tcx>,
diff --git a/compiler/rustc_trait_selection/src/solve/canonicalize.rs b/compiler/rustc_trait_selection/src/solve/canonicalize.rs
index 88771f90756..a9d182abfb7 100644
--- a/compiler/rustc_trait_selection/src/solve/canonicalize.rs
+++ b/compiler/rustc_trait_selection/src/solve/canonicalize.rs
@@ -207,23 +207,18 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'_, 'tcx> {
         t
     }
 
-    fn fold_region(&mut self, mut r: ty::Region<'tcx>) -> ty::Region<'tcx> {
-        match self.canonicalize_mode {
-            CanonicalizeMode::Input => {
-                // Don't resolve infer vars in input, since it affects
-                // caching and may cause trait selection bugs which rely
-                // on regions to be equal.
-            }
-            CanonicalizeMode::Response { .. } => {
-                if let ty::ReVar(vid) = *r {
-                    r = self
-                        .infcx
-                        .inner
-                        .borrow_mut()
-                        .unwrap_region_constraints()
-                        .opportunistic_resolve_var(self.infcx.tcx, vid);
-                }
-            }
+    fn fold_region(&mut self, r: ty::Region<'tcx>) -> ty::Region<'tcx> {
+        if let ty::ReVar(vid) = *r {
+            let resolved_region = self
+                .infcx
+                .inner
+                .borrow_mut()
+                .unwrap_region_constraints()
+                .opportunistic_resolve_var(self.infcx.tcx, vid);
+            assert_eq!(
+                r, resolved_region,
+                "region var should have been resolved, {r} -> {resolved_region}"
+            );
         }
 
         let kind = match *r {
@@ -278,38 +273,22 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'_, 'tcx> {
         ty::Region::new_late_bound(self.interner(), self.binder_index, br)
     }
 
-    fn fold_ty(&mut self, mut t: Ty<'tcx>) -> Ty<'tcx> {
+    fn fold_ty(&mut self, t: Ty<'tcx>) -> Ty<'tcx> {
         let kind = match *t.kind() {
-            ty::Infer(ty::TyVar(mut vid)) => {
-                // We need to canonicalize the *root* of our ty var.
-                // This is so that our canonical response correctly reflects
-                // any equated inference vars correctly!
-                let root_vid = self.infcx.root_var(vid);
-                if root_vid != vid {
-                    t = Ty::new_var(self.infcx.tcx, root_vid);
-                    vid = root_vid;
-                }
-
-                match self.infcx.probe_ty_var(vid) {
-                    Ok(t) => return self.fold_ty(t),
-                    Err(ui) => CanonicalVarKind::Ty(CanonicalTyVarKind::General(ui)),
-                }
+            ty::Infer(ty::TyVar(vid)) => {
+                assert_eq!(self.infcx.root_var(vid), vid, "ty vid should have been resolved");
+                let Err(ui) = self.infcx.probe_ty_var(vid) else {
+                    bug!("ty var should have been resolved: {t}");
+                };
+                CanonicalVarKind::Ty(CanonicalTyVarKind::General(ui))
             }
             ty::Infer(ty::IntVar(vid)) => {
-                let nt = self.infcx.opportunistic_resolve_int_var(vid);
-                if nt != t {
-                    return self.fold_ty(nt);
-                } else {
-                    CanonicalVarKind::Ty(CanonicalTyVarKind::Int)
-                }
+                assert_eq!(self.infcx.opportunistic_resolve_int_var(vid), t);
+                CanonicalVarKind::Ty(CanonicalTyVarKind::Int)
             }
             ty::Infer(ty::FloatVar(vid)) => {
-                let nt = self.infcx.opportunistic_resolve_float_var(vid);
-                if nt != t {
-                    return self.fold_ty(nt);
-                } else {
-                    CanonicalVarKind::Ty(CanonicalTyVarKind::Float)
-                }
+                assert_eq!(self.infcx.opportunistic_resolve_float_var(vid), t);
+                CanonicalVarKind::Ty(CanonicalTyVarKind::Float)
             }
             ty::Infer(ty::FreshTy(_) | ty::FreshIntTy(_) | ty::FreshFloatTy(_)) => {
                 bug!("fresh var during canonicalization: {t:?}")
@@ -372,22 +351,19 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'_, 'tcx> {
         Ty::new_bound(self.infcx.tcx, self.binder_index, bt)
     }
 
-    fn fold_const(&mut self, mut c: ty::Const<'tcx>) -> ty::Const<'tcx> {
+    fn fold_const(&mut self, c: ty::Const<'tcx>) -> ty::Const<'tcx> {
         let kind = match c.kind() {
-            ty::ConstKind::Infer(ty::InferConst::Var(mut vid)) => {
-                // We need to canonicalize the *root* of our const var.
-                // This is so that our canonical response correctly reflects
-                // any equated inference vars correctly!
-                let root_vid = self.infcx.root_const_var(vid);
-                if root_vid != vid {
-                    c = ty::Const::new_var(self.infcx.tcx, root_vid, c.ty());
-                    vid = root_vid;
-                }
-
-                match self.infcx.probe_const_var(vid) {
-                    Ok(c) => return self.fold_const(c),
-                    Err(universe) => CanonicalVarKind::Const(universe, c.ty()),
-                }
+            ty::ConstKind::Infer(ty::InferConst::Var(vid)) => {
+                assert_eq!(
+                    self.infcx.root_const_var(vid),
+                    vid,
+                    "const var should have been resolved"
+                );
+                let Err(ui) = self.infcx.probe_const_var(vid) else {
+                    bug!("const var should have been resolved");
+                };
+                // FIXME: we should fold this ty eventually
+                CanonicalVarKind::Const(ui, c.ty())
             }
             ty::ConstKind::Infer(ty::InferConst::Fresh(_)) => {
                 bug!("fresh var during canonicalization: {c:?}")
diff --git a/compiler/rustc_trait_selection/src/solve/eval_ctxt/canonical.rs b/compiler/rustc_trait_selection/src/solve/eval_ctxt/canonical.rs
index 749568cd80d..84ca555a5de 100644
--- a/compiler/rustc_trait_selection/src/solve/eval_ctxt/canonical.rs
+++ b/compiler/rustc_trait_selection/src/solve/eval_ctxt/canonical.rs
@@ -16,11 +16,15 @@ use rustc_index::IndexVec;
 use rustc_infer::infer::canonical::query_response::make_query_region_constraints;
 use rustc_infer::infer::canonical::CanonicalVarValues;
 use rustc_infer::infer::canonical::{CanonicalExt, QueryRegionConstraints};
+use rustc_infer::infer::InferCtxt;
 use rustc_middle::traits::query::NoSolution;
 use rustc_middle::traits::solve::{
- ExternalConstraintsData, MaybeCause, PredefinedOpaquesData, QueryInput,
+    ExternalConstraintsData, MaybeCause, PredefinedOpaquesData, QueryInput,
+};
+use rustc_middle::ty::{
+    self, BoundVar, GenericArgKind, Ty, TyCtxt, TypeFoldable, TypeFolder, TypeSuperFoldable,
+    TypeVisitableExt,
 };
-use rustc_middle::ty::{self, BoundVar, GenericArgKind, Ty, TyCtxt, TypeFoldable};
 use rustc_span::DUMMY_SP;
 use std::iter;
 use std::ops::Deref;
@@ -32,6 +36,10 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         &self,
         goal: Goal<'tcx, T>,
     ) -> (Vec<ty::GenericArg<'tcx>>, CanonicalInput<'tcx, T>) {
+        let opaque_types = self.infcx.clone_opaque_types_for_query_response();
+        let (goal, opaque_types) =
+            (goal, opaque_types).fold_with(&mut EagerResolver { infcx: self.infcx });
+
         let mut orig_values = Default::default();
         let canonical_goal = Canonicalizer::canonicalize(
             self.infcx,
@@ -40,11 +48,9 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             QueryInput {
                 goal,
                 anchor: self.infcx.defining_use_anchor,
-                predefined_opaques_in_body: self.tcx().mk_predefined_opaques_in_body(
-                    PredefinedOpaquesData {
-                        opaque_types: self.infcx.clone_opaque_types_for_query_response(),
-                    },
-                ),
+                predefined_opaques_in_body: self
+                    .tcx()
+                    .mk_predefined_opaques_in_body(PredefinedOpaquesData { opaque_types }),
             },
         );
         (orig_values, canonical_goal)
@@ -69,6 +75,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             previous call to `try_evaluate_added_goals!`"
         );
 
+        let certainty = certainty.unify_with(goals_certainty);
         if let Certainty::OVERFLOW = certainty {
             // If we have overflow, it's probable that we're substituting a type
             // into itself infinitely and any partial substitutions in the query
@@ -84,10 +91,12 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             return Ok(self.make_ambiguous_response_no_constraints(MaybeCause::Overflow));
         }
 
-        let certainty = certainty.unify_with(goals_certainty);
         let var_values = self.var_values;
         let external_constraints = self.compute_external_query_constraints()?;
 
+        let (var_values, external_constraints) =
+            (var_values, external_constraints).fold_with(&mut EagerResolver { infcx: self.infcx });
+
         let canonical = Canonicalizer::canonicalize(
             self.infcx,
             CanonicalizeMode::Response { max_input_universe: self.max_input_universe },
@@ -334,3 +343,65 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         Ok(())
     }
 }
+
+/// Resolves ty, region, and const vars to their inferred values or their root vars.
+struct EagerResolver<'a, 'tcx> {
+    infcx: &'a InferCtxt<'tcx>,
+}
+
+impl<'tcx> TypeFolder<TyCtxt<'tcx>> for EagerResolver<'_, 'tcx> {
+    fn interner(&self) -> TyCtxt<'tcx> {
+        self.infcx.tcx
+    }
+
+    fn fold_ty(&mut self, t: Ty<'tcx>) -> Ty<'tcx> {
+        match *t.kind() {
+            ty::Infer(ty::TyVar(vid)) => match self.infcx.probe_ty_var(vid) {
+                Ok(t) => t.fold_with(self),
+                Err(_) => Ty::new_var(self.infcx.tcx, self.infcx.root_var(vid)),
+            },
+            ty::Infer(ty::IntVar(vid)) => self.infcx.opportunistic_resolve_int_var(vid),
+            ty::Infer(ty::FloatVar(vid)) => self.infcx.opportunistic_resolve_float_var(vid),
+            _ => {
+                if t.has_infer() {
+                    t.super_fold_with(self)
+                } else {
+                    t
+                }
+            }
+        }
+    }
+
+    fn fold_region(&mut self, r: ty::Region<'tcx>) -> ty::Region<'tcx> {
+        match *r {
+            ty::ReVar(vid) => self
+                .infcx
+                .inner
+                .borrow_mut()
+                .unwrap_region_constraints()
+                .opportunistic_resolve_var(self.infcx.tcx, vid),
+            _ => r,
+        }
+    }
+
+    fn fold_const(&mut self, c: ty::Const<'tcx>) -> ty::Const<'tcx> {
+        match c.kind() {
+            ty::ConstKind::Infer(ty::InferConst::Var(vid)) => {
+                // FIXME: we need to fold the ty too, I think.
+                match self.infcx.probe_const_var(vid) {
+                    Ok(c) => c.fold_with(self),
+                    Err(_) => {
+                        ty::Const::new_var(self.infcx.tcx, self.infcx.root_const_var(vid), c.ty())
+                    }
+                }
+            }
+            _ => {
+                if c.has_infer() {
+                    c.super_fold_with(self)
+                } else {
+                    c
+                }
+            }
+        }
+    }
+}