Clean up and speed up inference variable resolving code

compiler/rustc_hir_analysis/src/check/compare_impl_item.rs

@@ -1297,6 +1297,7 @@ fn check_region_late_boundedness<'tcx>(
                 .borrow_mut()
                 .unwrap_region_constraints()
                 .shallow_resolve_region_var(tcx, vid)
+                .reuse_if_unchanged(r)
             && let ty::ReLateParam(ty::LateParamRegion {
                 kind: ty::LateParamRegionKind::Named(trait_param_def_id),
                 ..
@@ -1322,6 +1323,7 @@ fn check_region_late_boundedness<'tcx>(
                 .borrow_mut()
                 .unwrap_region_constraints()
                 .shallow_resolve_region_var(tcx, vid)
+                .reuse_if_unchanged(r)
             && let ty::ReLateParam(ty::LateParamRegion {
                 kind: ty::LateParamRegionKind::Named(impl_param_def_id),
                 ..

compiler/rustc_hir_typeck/src/closure.rs

@@ -286,7 +286,9 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     }
 
     /// Given the expected type, figures out what it can about this closure we
-    /// are about to type check:
+    /// are about to type check.
+    ///
+    /// WARNING: `expected_ty` must be resolved, to ensure that tyvars refer to root vids.
     #[instrument(skip(self), level = "debug", ret)]
     fn deduce_closure_signature(
         &self,
@@ -313,13 +315,17 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     .and_then(|did| self.tcx.fn_trait_kind_from_def_id(did));
                 (sig, kind)
             }
-            ty::Infer(ty::TyVar(vid)) => self.deduce_closure_signature_from_predicates(
-                Ty::new_var(self.tcx, self.root_var(vid)),
-                closure_kind,
-                self.obligations_for_self_ty(vid, UseSubtyping::No)
-                    .into_iter()
-                    .filter_map(|obl| Some((obl.predicate.as_clause()?, obl.cause.span))),
-            ),
+            ty::Infer(ty::TyVar(vid)) => {
+                // assert that the precondition (documented in the doc comments) is maintained.
+                debug_assert_eq!(self.root_ty_var(vid), vid);
+                self.deduce_closure_signature_from_predicates(
+                    Ty::new_var(self.tcx, vid),
+                    closure_kind,
+                    self.obligations_for_self_ty(vid, UseSubtyping::No)
+                        .into_iter()
+                        .filter_map(|obl| Some((obl.predicate.as_clause()?, obl.cause.span))),
+                )
+            }
             ty::FnPtr(sig_tys, hdr) => match closure_kind {
                 hir::ClosureKind::Closure => {
                     let expected_sig = ExpectedSig { cause_span: None, sig: sig_tys.with(hdr) };

compiler/rustc_hir_typeck/src/fallback.rs

@@ -315,12 +315,12 @@ impl<'tcx> FnCtxt<'_, 'tcx> {
 
     /// If `ty` is an unresolved type variable, returns its root vid.
     fn root_vid(&self, ty: Ty<'tcx>) -> Option<ty::TyVid> {
-        Some(self.root_var(self.shallow_resolve(ty).ty_vid()?))
+        Some(self.shallow_resolve(ty).ty_vid()?)
     }
 
     /// If `ty` is an unresolved float type variable, returns its root vid.
     pub(crate) fn root_float_vid(&self, ty: Ty<'tcx>) -> Option<ty::FloatVid> {
-        Some(self.root_float_var(self.shallow_resolve(ty).float_vid()?))
+        Some(self.shallow_resolve(ty).float_vid()?)
     }
 
     /// Given a set of diverging vids and coercions, walk the HIR to gather a

compiler/rustc_hir_typeck/src/fn_ctxt/inspect_obligations.rs

@@ -95,7 +95,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
 
         match *ty.kind() {
             ty::Infer(ty::TyVar(found_vid)) => match subtyping {
-                UseSubtyping::No => self.root_var(expected_vid) == self.root_var(found_vid),
+                UseSubtyping::No => self.root_ty_var(expected_vid) == found_vid,
                 UseSubtyping::Yes => {
                     self.sub_unification_table_root_var(expected_vid)
                         == self.sub_unification_table_root_var(found_vid)

compiler/rustc_infer/src/infer/canonical/canonicalizer.rs

@@ -165,7 +165,8 @@ impl CanonicalizeMode for CanonicalizeQueryResponse {
                 .inner
                 .borrow_mut()
                 .unwrap_region_constraints()
-                .shallow_resolve_region_var(canonicalizer.tcx, vid);
+                .shallow_resolve_region_var(canonicalizer.tcx, vid)
+                .reuse_if_unchanged(r);
             debug!(
                 "canonical: region var found with vid {vid:?}, \
                      opportunistically resolved to {r:?}",
@@ -335,7 +336,7 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
                 // 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.unwrap().root_var(vid);
+                let root_vid = self.infcx.unwrap().root_ty_var(vid);
                 if root_vid != vid {
                     t = Ty::new_var(self.tcx, root_vid);
                     vid = root_vid;
@@ -363,15 +364,15 @@ impl<'cx, 'tcx> TypeFolder<TyCtxt<'tcx>> for Canonicalizer<'cx, 'tcx> {
             }
 
             ty::Infer(ty::IntVar(vid)) => {
-                let nt = self.infcx.unwrap().shallow_resolve_int_var(vid);
+                let nt = self.infcx.unwrap().shallow_resolve_int_var(vid).reuse_if_unchanged(t);
                 if nt != t {
                     return self.fold_ty(nt);
                 } else {
                     self.canonicalize_ty_var(CanonicalVarKind::Int, t)
                 }
             }
             ty::Infer(ty::FloatVar(vid)) => {
-                let nt = self.infcx.unwrap().shallow_resolve_float_var(vid);
+                let nt = self.infcx.unwrap().shallow_resolve_float_var(vid).reuse_if_unchanged(t);
                 if nt != t {
                     return self.fold_ty(nt);
                 } else {

compiler/rustc_infer/src/infer/context.rs

@@ -97,7 +97,7 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
     }
 
     fn root_ty_var(&self, var: ty::TyVid) -> ty::TyVid {
-        self.root_var(var)
+        self.root_ty_var(var)
     }
 
     fn sub_unification_table_root_var(&self, var: ty::TyVid) -> ty::TyVid {
@@ -118,23 +118,32 @@ impl<'tcx> rustc_type_ir::InferCtxtLike for InferCtxt<'tcx> {
         self.root_const_var(var)
     }
 
-    fn shallow_resolve_ty_var(&self, vid: ty::TyVid) -> Ty<'tcx> {
+    fn shallow_resolve_ty_var(&self, vid: ty::TyVid) -> ty::MaybeResolved<'tcx, ty::TyVid> {
         self.shallow_resolve_ty_var(vid)
     }
 
-    fn shallow_resolve_int_var(&self, vid: ty::IntVid) -> Ty<'tcx> {
+    fn shallow_resolve_int_var(&self, vid: ty::IntVid) -> ty::MaybeResolved<'tcx, ty::IntVid> {
         self.shallow_resolve_int_var(vid)
     }
 
-    fn shallow_resolve_float_var(&self, vid: ty::FloatVid) -> Ty<'tcx> {
+    fn shallow_resolve_float_var(
+        &self,
+        vid: ty::FloatVid,
+    ) -> ty::MaybeResolved<'tcx, ty::FloatVid> {
         self.shallow_resolve_float_var(vid)
     }
 
-    fn shallow_resolve_const_var(&self, vid: ty::ConstVid) -> ty::Const<'tcx> {
+    fn shallow_resolve_const_var(
+        &self,
+        vid: ty::ConstVid,
+    ) -> ty::MaybeResolved<'tcx, ty::ConstVid> {
         self.shallow_resolve_const_var(vid)
     }
 
-    fn shallow_resolve_region_var(&self, vid: ty::RegionVid) -> ty::Region<'tcx> {
+    fn shallow_resolve_region_var(
+        &self,
+        vid: ty::RegionVid,
+    ) -> ty::MaybeResolved<'tcx, ty::RegionVid> {
         self.inner
             .borrow_mut()
             .unwrap_region_constraints()
@@ -449,7 +458,7 @@ impl<'a, 'tcx> ty::TypeFolder<TyCtxt<'tcx>> for LowerUniverseFolder<'a, 'tcx> {
 
         let folded = match t.kind() {
             ty::Infer(ty::TyVar(vid)) => {
-                let vid = self.infcx.root_var(*vid);
+                let vid = self.infcx.root_ty_var(*vid);
                 let probe = self.infcx.inner.borrow_mut().type_variables().probe(vid);
                 match probe {
                     TypeVariableValue::Known { value: u } => u.super_fold_with(self),
@@ -481,8 +490,8 @@ impl<'a, 'tcx> ty::TypeFolder<TyCtxt<'tcx>> for LowerUniverseFolder<'a, 'tcx> {
 
         match c.kind() {
             ty::ConstKind::Infer(ty::InferConst::Var(vid)) => {
-                let vid = self.infcx.root_const_var(vid);
-                let universe = match self.infcx.try_resolve_const_var(vid) {
+                let (res, vid) = self.infcx.try_resolve_const_var_with_root(vid);
+                let universe = match res {
                     Ok(value) => return value.fold_with(self),
                     Err(universe) => universe,
                 };