Fix canonicalizer bug for int/float vars too

2 files changed, 44 insertions, 9 deletions

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