fixup various sites that use root_{ty,const}_var that can skip it after a shallow_resolve

Author: jdonszelmann

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

@@ -399,12 +399,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

@@ -334,7 +334,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;

compiler/rustc_infer/src/infer/context.rs

@@ -103,7 +103,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 {
@@ -461,7 +461,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),
@@ -493,8 +493,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,
                 };

compiler/rustc_infer/src/infer/mod.rs

@@ -1222,8 +1222,29 @@ impl<'tcx> InferCtxt<'tcx> {
         self.deeply_resolve_ignoring_regions(t).to_string()
     }
 
-    /// If `TyVar(vid)` resolves to a type, return that type. Else, return the
-    /// universe index of `TyVar(vid)`.
+    /// If `TyVar(vid)` resolves to a type, return that type.
+    /// Else, return the universe index of `TyVar(vid)`.
+    ///
+    /// Also return the root `TyVid` of `vid`.
+    /// This is more efficient than calling [`try_resolve_ty_var`](Self::try_resolve_ty_var)
+    /// followed by [`root_ty_var`](Self::root_ty_var).
+    pub fn try_resolve_ty_var_with_root(
+        &self,
+        vid: TyVid,
+    ) -> (Result<Ty<'tcx>, ty::UniverseIndex>, TyVid) {
+        let (root, value) = self.inner.borrow_mut().type_variables().probe_with_root_vid(vid);
+
+        (
+            match value {
+                TypeVariableValue::Known { value } => Ok(self.shallow_resolve_non_recursive(value)),
+                TypeVariableValue::Unknown { universe } => Err(universe),
+            },
+            root,
+        )
+    }
+
+    /// If `TyVar(vid)` resolves to a type, return that type.
+    /// Else, return the universe index of `TyVar(vid)`.
     pub fn try_resolve_ty_var(&self, vid: TyVid) -> Result<Ty<'tcx>, ty::UniverseIndex> {
         let value = self.inner.borrow_mut().type_variables().probe(vid);
 
@@ -1235,12 +1256,8 @@ impl<'tcx> InferCtxt<'tcx> {
 
     /// If `vid` resolves to a type, return that type. Otherwise return the root variable id for `vid`.
     pub fn shallow_resolve_ty_var_or_get_root(&self, vid: TyVid) -> Result<Ty<'tcx>, TyVid> {
-        let (root, value) = self.inner.borrow_mut().type_variables().probe_with_root_vid(vid);
-
-        match value {
-            TypeVariableValue::Known { value } => Ok(self.shallow_resolve_non_recursive(value)),
-            TypeVariableValue::Unknown { universe: _ } => Err(root),
-        }
+        let (res, root) = self.try_resolve_ty_var_with_root(vid);
+        res.map_err(|_| root)
     }
 
     /// Resolve a type variable to a type, if known.
@@ -1441,7 +1458,7 @@ impl<'tcx> InferCtxt<'tcx> {
         }
     }
 
-    pub fn root_var(&self, var: ty::TyVid) -> ty::TyVid {
+    pub fn root_ty_var(&self, var: ty::TyVid) -> ty::TyVid {
         self.inner.borrow_mut().type_variables().root_var(var)
     }
 
@@ -1524,6 +1541,29 @@ impl<'tcx> InferCtxt<'tcx> {
         value.fold_with(&mut r)
     }
 
+    /// If `ConstVar(vid)` resolves to a const, return that const.
+    /// Else, return the universe index of `ConstVar(vid)`.
+    ///
+    /// Also return the root `ConstVid` of `vid`.
+    /// This is more efficient than calling [`try_resolve_const_var`](Self::try_resolve_const_var)
+    /// followed by [`root_const_var`](Self::root_const_var).
+    pub fn try_resolve_const_var_with_root(
+        &self,
+        vid: ty::ConstVid,
+    ) -> (Result<ty::Const<'tcx>, ty::UniverseIndex>, ty::ConstVid) {
+        let (root, value) =
+            self.inner.borrow_mut().const_unification_table().inlined_probe_key_value(vid);
+        (
+            match value {
+                ConstVariableValue::Known { value } => Ok(value),
+                ConstVariableValue::Unknown { origin: _, universe } => Err(universe),
+            },
+            root.vid,
+        )
+    }
+
+    /// If `ConstVar(vid)` resolves to a const, return that const.
+    /// Else, return the universe index of `ConstVar(vid)`.
     pub fn try_resolve_const_var(
         &self,
         vid: ty::ConstVid,

compiler/rustc_infer/src/infer/relate/generalize.rs

@@ -287,12 +287,13 @@ impl<'tcx> InferCtxt<'tcx> {
         assert!(!source_term.has_escaping_bound_vars());
         let (for_universe, root_vid) = match target_vid {
             TermVid::Ty(ty_vid) => {
-                (self.try_resolve_ty_var(ty_vid).unwrap_err(), TermVid::Ty(self.root_var(ty_vid)))
+                let (res, root) = self.try_resolve_ty_var_with_root(ty_vid);
+                (res.unwrap_err(), TermVid::Ty(root))
+            }
+            TermVid::Const(ct_vid) => {
+                let (res, root) = self.try_resolve_const_var_with_root(ct_vid);
+                (res.unwrap_err(), TermVid::Const(root))
             }
-            TermVid::Const(ct_vid) => (
-                self.try_resolve_const_var(ct_vid).unwrap_err(),
-                TermVid::Const(self.inner.borrow_mut().const_unification_table().find(ct_vid).vid),
-            ),
         };
 
         let mut generalizer = Generalizer {