From 1b6d6f92d1980d7f987fa0fe53cea8d8fb3209eb Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Wed, 8 Mar 2023 22:10:24 +0000
Subject: [PATCH] Fix canonicalizer bug for int/float vars too

---
 .../src/infer/canonical/canonicalizer.rs      | 31 +++++++++++++------
 compiler/rustc_infer/src/infer/mod.rs         | 22 +++++++++++++
 .../src/solve/canonical/canonicalize.rs       |  8 ++---
 .../canonical-int-var-eq-in-response.rs       | 21 +++++++++++++
 4 files changed, 69 insertions(+), 13 deletions(-)
 create mode 100644 tests/ui/traits/new-solver/canonical-int-var-eq-in-response.rs

diff --git a/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs b/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs
index a43330e2d4e..8ac82653c0e 100644
--- a/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs
+++ b/compiler/rustc_infer/src/infer/canonical/canonicalizer.rs
@@ -411,15 +411,28 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
                 }
             }
 
-            ty::Infer(ty::IntVar(_)) => self.canonicalize_ty_var(
-                CanonicalVarInfo { kind: CanonicalVarKind::Ty(CanonicalTyVarKind::Int) },
-                t,
-            ),
-
-            ty::Infer(ty::FloatVar(_)) => self.canonicalize_ty_var(
-                CanonicalVarInfo { kind: CanonicalVarKind::Ty(CanonicalTyVarKind::Float) },
-                t,
-            ),
+            ty::Infer(ty::IntVar(vid)) => {
+                let nt = self.infcx.opportunistic_resolve_int_var(vid);
+                if nt != t {
+                    return self.fold_ty(nt);
+                } else {
+                    self.canonicalize_ty_var(
+                        CanonicalVarInfo { kind: CanonicalVarKind::Ty(CanonicalTyVarKind::Int) },
+                        t,
+                    )
+                }
+            }
+            ty::Infer(ty::FloatVar(vid)) => {
+                let nt = self.infcx.opportunistic_resolve_float_var(vid);
+                if nt != t {
+                    return self.fold_ty(nt);
+                } else {
+                    self.canonicalize_ty_var(
+                        CanonicalVarInfo { kind: CanonicalVarKind::Ty(CanonicalTyVarKind::Float) },
+                        t,
+                    )
+                }
+            }
 
             ty::Infer(ty::FreshTy(_) | ty::FreshIntTy(_) | ty::FreshFloatTy(_)) => {
                 bug!("encountered a fresh type during canonicalization")
diff --git a/compiler/rustc_infer/src/infer/mod.rs b/compiler/rustc_infer/src/infer/mod.rs
index 4a2a5557313..4a834957959 100644
--- a/compiler/rustc_infer/src/infer/mod.rs
+++ b/compiler/rustc_infer/src/infer/mod.rs
@@ -1363,6 +1363,28 @@ impl<'tcx> InferCtxt<'tcx> {
         self.inner.borrow_mut().const_unification_table().find(var)
     }
 
+    /// Resolves an int var to a rigid int type, if it was constrained to one,
+    /// or else the root int var in the unification table.
+    pub fn opportunistic_resolve_int_var(&self, vid: ty::IntVid) -> Ty<'tcx> {
+        let mut inner = self.inner.borrow_mut();
+        if let Some(value) = inner.int_unification_table().probe_value(vid) {
+            value.to_type(self.tcx)
+        } else {
+            self.tcx.mk_int_var(inner.int_unification_table().find(vid))
+        }
+    }
+
+    /// Resolves a float var to a rigid int type, if it was constrained to one,
+    /// or else the root float var in the unification table.
+    pub fn opportunistic_resolve_float_var(&self, vid: ty::FloatVid) -> Ty<'tcx> {
+        let mut inner = self.inner.borrow_mut();
+        if let Some(value) = inner.float_unification_table().probe_value(vid) {
+            value.to_type(self.tcx)
+        } else {
+            self.tcx.mk_float_var(inner.float_unification_table().find(vid))
+        }
+    }
+
     /// Where possible, replaces type/const variables in
     /// `value` with their final value. Note that region variables
     /// are unaffected. If a type/const variable has not been unified, it
diff --git a/compiler/rustc_trait_selection/src/solve/canonical/canonicalize.rs b/compiler/rustc_trait_selection/src/solve/canonical/canonicalize.rs
index 981a8f45e45..7ee4f332306 100644
--- a/compiler/rustc_trait_selection/src/solve/canonical/canonicalize.rs
+++ b/compiler/rustc_trait_selection/src/solve/canonical/canonicalize.rs
@@ -278,16 +278,16 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'_, 'tcx> {
                     Err(ui) => CanonicalVarKind::Ty(CanonicalTyVarKind::General(ui)),
                 }
             }
-            ty::Infer(ty::IntVar(_)) => {
-                let nt = self.infcx.shallow_resolve(t);
+            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)
                 }
             }
-            ty::Infer(ty::FloatVar(_)) => {
-                let nt = self.infcx.shallow_resolve(t);
+            ty::Infer(ty::FloatVar(vid)) => {
+                let nt = self.infcx.opportunistic_resolve_float_var(vid);
                 if nt != t {
                     return self.fold_ty(nt);
                 } else {
diff --git a/tests/ui/traits/new-solver/canonical-int-var-eq-in-response.rs b/tests/ui/traits/new-solver/canonical-int-var-eq-in-response.rs
new file mode 100644
index 00000000000..3f7316a2279
--- /dev/null
+++ b/tests/ui/traits/new-solver/canonical-int-var-eq-in-response.rs
@@ -0,0 +1,21 @@
+// check-pass
+
+trait Mirror {
+    type Assoc;
+}
+
+impl<T> Mirror for T {
+    type Assoc = T;
+}
+
+trait Test {}
+impl Test for i64 {}
+impl Test for u64 {}
+
+fn mirror_me<T: Mirror>(t: T, s: <T as Mirror>::Assoc) where <T as Mirror>::Assoc: Test {}
+
+fn main() {
+    let mut x = 0;
+    mirror_me(x, 1);
+    x = 1i64;
+}