Shallow resolve ty and const vars to their root vars, attempt 2

Cargo.lock

@@ -1345,9 +1345,9 @@ checksum = "edd0f118536f44f5ccd48bcb8b111bdc3de888b58c74639dfb034a357d0f206d"
 
 [[package]]
 name = "ena"
-version = "0.14.3"
+version = "0.14.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3d248bdd43ce613d87415282f69b9bb99d947d290b10962dd6c56233312c2ad5"
+checksum = "eabffdaee24bd1bf95c5ef7cec31260444317e72ea56c4c91750e8b7ee58d5f1"
 dependencies = [
  "log",
 ]

compiler/rustc_data_structures/Cargo.toml

@@ -9,7 +9,7 @@ arrayvec = { version = "0.7", default-features = false }
 bitflags = "2.4.1"
 either = "1.0"
 elsa = "1.11.0"
-ena = "0.14.3"
+ena = "0.14.4"
 indexmap = "2.14.0"
 jobserver_crate = { version = "0.1.28", package = "jobserver" }
 measureme = "12.0.1"

compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs

@@ -15,9 +15,12 @@ use rustc_hir::{Expr, ExprKind, FnRetTy, HirId, LangItem, Node, QPath, is_range_
 use rustc_hir_analysis::check::potentially_plural_count;
 use rustc_hir_analysis::hir_ty_lowering::{HirTyLowerer, ResolvedStructPath};
 use rustc_index::IndexVec;
-use rustc_infer::infer::{BoundRegionConversionTime, DefineOpaqueTypes, InferOk, TypeTrace};
+use rustc_infer::infer::{
+    BoundRegionConversionTime, DefineOpaqueTypes, InferOk, TypeTrace, relate,
+};
 use rustc_middle::ty::adjustment::AllowTwoPhase;
-use rustc_middle::ty::error::TypeError;
+use rustc_middle::ty::error::{ExpectedFound, TypeError};
+use rustc_middle::ty::relate::{Relate, RelateResult, TypeRelation};
 use rustc_middle::ty::{self, IsSuggestable, Ty, TyCtxt, TypeVisitableExt, Unnormalized};
 use rustc_middle::{bug, span_bug};
 use rustc_session::Session;
@@ -278,12 +281,29 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                         formal_input_tys
                             .iter()
                             .map(|&ty| self.resolve_vars_if_possible(ty))
-                            .collect(),
+                            .collect::<Vec<_>>(),
                     ))
                 })
                 .ok()
             })
-            .unwrap_or_default();
+            .unwrap_or_default()
+            .map(|expected_input_tys| {
+                expected_input_tys
+                    .into_iter()
+                    .zip(formal_input_tys)
+                    // if the expected input type is structurally equal to the formal input type,
+                    // i.e. we've only changed some inference variables around, keep the formal
+                    // input ty as the expected input ty.
+                    .map(|(expected_input_ty, formal_input_ty)| {
+                        if same_type_modulo_vars(tcx, expected_input_ty, *formal_input_ty) {
+                            // if they're the same, fall back to the formal input type
+                            *formal_input_ty
+                        } else {
+                            expected_input_ty
+                        }
+                    })
+                    .collect()
+            });
 
         let mut err_code = E0061;
 
@@ -3519,3 +3539,98 @@ enum SuggestionText {
     Reorder,
     DidYouMean,
 }
+
+fn same_type_modulo_vars<'tcx>(tcx: TyCtxt<'tcx>, a: Ty<'tcx>, b: Ty<'tcx>) -> bool {
+    struct SameModuloVars<'tcx> {
+        tcx: TyCtxt<'tcx>,
+    }
+    impl<'tcx> TypeRelation<TyCtxt<'tcx>> for SameModuloVars<'tcx> {
+        fn cx(&self) -> TyCtxt<'tcx> {
+            self.tcx
+        }
+
+        fn relate_ty_args(
+            &mut self,
+            a_ty: Ty<'tcx>,
+            _b_ty: Ty<'tcx>,
+            _ty_def_id: DefId,
+            a_args: ty::GenericArgsRef<'tcx>,
+            b_args: ty::GenericArgsRef<'tcx>,
+            _mk: impl FnOnce(ty::GenericArgsRef<'tcx>) -> Ty<'tcx>,
+        ) -> RelateResult<'tcx, Ty<'tcx>> {
+            relate::relate_args_invariantly(self, a_args, b_args)?;
+            Ok(a_ty)
+        }
+
+        fn relate_with_variance<T: Relate<TyCtxt<'tcx>>>(
+            &mut self,
+            _variance: ty::Variance,
+            _info: ty::VarianceDiagInfo<TyCtxt<'tcx>>,
+            a: T,
+            b: T,
+        ) -> RelateResult<'tcx, T> {
+            self.relate(a, b)
+        }
+
+        fn tys(&mut self, a: Ty<'tcx>, b: Ty<'tcx>) -> RelateResult<'tcx, Ty<'tcx>> {
+            if a == b {
+                return Ok(a);
+            }
+
+            match (a.kind(), b.kind()) {
+                (&ty::Infer(ty::InferTy::TyVar(_)), &ty::Infer(ty::InferTy::TyVar(_)))
+                | (&ty::Infer(ty::InferTy::FloatVar(_)), &ty::Infer(ty::InferTy::FloatVar(_)))
+                | (&ty::Infer(ty::InferTy::IntVar(_)), &ty::Infer(ty::InferTy::IntVar(_))) => Ok(a),
+                (&ty::Infer(_), _) | (_, &ty::Infer(_)) => Err(TypeError::Mismatch),
+                (&ty::Error(guar), _) | (_, &ty::Error(guar)) => Ok(Ty::new_error(self.cx(), guar)),
+                _ => relate::structurally_relate_tys(self, a, b),
+            }
+        }
+
+        fn regions(
+            &mut self,
+            a: ty::Region<'tcx>,
+            _b: ty::Region<'tcx>,
+        ) -> RelateResult<'tcx, ty::Region<'tcx>> {
+            Ok(a)
+        }
+
+        fn consts(
+            &mut self,
+            mut a: ty::Const<'tcx>,
+            mut b: ty::Const<'tcx>,
+        ) -> RelateResult<'tcx, ty::Const<'tcx>> {
+            if a == b {
+                return Ok(a);
+            }
+
+            if self.tcx.features().generic_const_exprs() {
+                a = self.tcx.expand_abstract_consts(a);
+                b = self.tcx.expand_abstract_consts(b);
+            }
+
+            match (a.kind(), b.kind()) {
+                (ty::ConstKind::Infer(_), ty::ConstKind::Infer(_)) => return Ok(a),
+                (ty::ConstKind::Infer(_), _) | (_, ty::ConstKind::Infer(_)) => {
+                    return Err(TypeError::ConstMismatch(ExpectedFound::new(a, b)));
+                }
+                _ => {}
+            }
+
+            relate::structurally_relate_consts(self, a, b)
+        }
+
+        fn binders<T>(
+            &mut self,
+            a: ty::Binder<'tcx, T>,
+            b: ty::Binder<'tcx, T>,
+        ) -> RelateResult<'tcx, ty::Binder<'tcx, T>>
+        where
+            T: Relate<TyCtxt<'tcx>>,
+        {
+            Ok(a.rebind(self.relate(a.skip_binder(), b.skip_binder())?))
+        }
+    }
+
+    SameModuloVars { tcx }.relate(a, b).is_ok()
+}

compiler/rustc_infer/src/infer/mod.rs

@@ -1224,22 +1224,45 @@ impl<'tcx> InferCtxt<'tcx> {
                     //
                     // Note: if these two lines are combined into one we get
                     // dynamic borrow errors on `self.inner`.
-                    let known = self.inner.borrow_mut().type_variables().probe(v).known();
-                    known.map_or(ty, |t| self.shallow_resolve(t))
+                    let (root_vid, value) =
+                        self.inner.borrow_mut().type_variables().probe_with_root_vid(v);
+                    value.known().map_or_else(
+                        || if root_vid == v { ty } else { Ty::new_var(self.tcx, root_vid) },
+                        |t| self.shallow_resolve(t),
+                    )
                 }
 
                 ty::IntVar(v) => {
-                    match self.inner.borrow_mut().int_unification_table().probe_value(v) {
+                    let (root, value) =
+                        self.inner.borrow_mut().int_unification_table().inlined_probe_key_value(v);
+                    match value {
                         ty::IntVarValue::IntType(ty) => Ty::new_int(self.tcx, ty),
                         ty::IntVarValue::UintType(ty) => Ty::new_uint(self.tcx, ty),
-                        ty::IntVarValue::Unknown => ty,
+                        ty::IntVarValue::Unknown => {
+                            if root == v {
+                                ty
+                            } else {
+                                Ty::new_int_var(self.tcx, root)
+                            }
+                        }
                     }
                 }
 
                 ty::FloatVar(v) => {
-                    match self.inner.borrow_mut().float_unification_table().probe_value(v) {
+                    let (root, value) = self
+                        .inner
+                        .borrow_mut()
+                        .float_unification_table()
+                        .inlined_probe_key_value(v);
+                    match value {
                         ty::FloatVarValue::Known(ty) => Ty::new_float(self.tcx, ty),
-                        ty::FloatVarValue::Unknown => ty,
+                        ty::FloatVarValue::Unknown => {
+                            if root == v {
+                                ty
+                            } else {
+                                Ty::new_float_var(self.tcx, root)
+                            }
+                        }
                     }
                 }
 
@@ -1253,13 +1276,16 @@ impl<'tcx> InferCtxt<'tcx> {
     pub fn shallow_resolve_const(&self, ct: ty::Const<'tcx>) -> ty::Const<'tcx> {
         match ct.kind() {
             ty::ConstKind::Infer(infer_ct) => match infer_ct {
-                InferConst::Var(vid) => self
-                    .inner
-                    .borrow_mut()
-                    .const_unification_table()
-                    .probe_value(vid)
-                    .known()
-                    .unwrap_or(ct),
+                InferConst::Var(vid) => {
+                    let (root, value) = self
+                        .inner
+                        .borrow_mut()
+                        .const_unification_table()
+                        .inlined_probe_key_value(vid);
+                    value.known().unwrap_or_else(|| {
+                        if root.vid == vid { ct } else { ty::Const::new_var(self.tcx, root.vid) }
+                    })
+                }
                 InferConst::Fresh(_) => ct,
             },
 
@@ -1284,6 +1310,13 @@ impl<'tcx> InferCtxt<'tcx> {
         self.inner.borrow_mut().type_variables().root_var(var)
     }
 
+    /// If `ty` is an unresolved type variable, returns its root vid.
+    pub fn root_vid(&self, ty: Ty<'tcx>) -> Option<ty::TyVid> {
+        let (root, value) =
+            self.inner.borrow_mut().type_variables().inlined_probe_with_vid(ty.ty_vid()?);
+        value.is_unknown().then_some(root)
+    }
+
     pub fn sub_unify_ty_vids_raw(&self, a: ty::TyVid, b: ty::TyVid) {
         self.inner.borrow_mut().type_variables().sub_unify(a, b);
     }

compiler/rustc_infer/src/infer/type_variable.rs

@@ -267,6 +267,25 @@ impl<'tcx> TypeVariableTable<'_, 'tcx> {
         self.eq_relations().inlined_probe_value(vid)
     }
 
+    /// Retrieves the type to which `vid` has been instantiated, if
+    /// any, along with the root `vid`.
+    pub(crate) fn probe_with_root_vid(
+        &mut self,
+        vid: ty::TyVid,
+    ) -> (ty::TyVid, TypeVariableValue<'tcx>) {
+        self.inlined_probe_with_vid(vid)
+    }
+
+    /// An always-inlined variant of `probe_with_root_vid`, for very hot call sites.
+    #[inline(always)]
+    pub(crate) fn inlined_probe_with_vid(
+        &mut self,
+        vid: ty::TyVid,
+    ) -> (ty::TyVid, TypeVariableValue<'tcx>) {
+        let (id, value) = self.eq_relations().inlined_probe_key_value(vid);
+        (id.vid, value)
+    }
+
     #[inline]
     fn eq_relations(&mut self) -> super::UnificationTable<'_, 'tcx, TyVidEqKey<'tcx>> {
         self.storage.eq_relations.with_log(self.undo_log)

compiler/rustc_type_ir/Cargo.toml

@@ -8,7 +8,7 @@ edition = "2024"
 arrayvec = { version = "0.7", default-features = false }
 bitflags = "2.4.1"
 derive-where = "1.6.1"
-ena = "0.14.3"
+ena = "0.14.4"
 indexmap = "2.0.0"
 rustc-hash = "2.0.0"
 rustc_abi = { path = "../rustc_abi", default-features = false }

tests/incremental/const-generics/issue-64087.rs

@@ -6,6 +6,4 @@ fn combinator<T, const S: usize>() -> [T; S] {}
 fn main() {
     combinator().into_iter();
     //[bfail1]~^ ERROR type annotations needed
-    //[bfail1]~| ERROR type annotations needed
-    //[bfail1]~| ERROR type annotations needed
 }

tests/ui/associated-inherent-types/inference-fail.stderr

@@ -2,7 +2,7 @@ error[E0282]: type annotations needed
   --> $DIR/inference-fail.rs:10:12
    |
 LL |     let _: S<_>::P = ();
-   |            ^^^^^^^ cannot infer type for type parameter `T`
+   |            ^^^^^^^ cannot infer type
 
 error: aborting due to 1 previous error
 

tests/ui/borrowck/index-mut-help2.stderr

@@ -25,10 +25,10 @@ LL |     map.insert(*****index, 23);
    |                ++++
 
 error[E0277]: the trait bound `&B: Borrow<&&&&B>` is not satisfied
-  --> $DIR/index-mut-help2.rs:95:9
+  --> $DIR/index-mut-help2.rs:95:5
    |
 LL |     map[index] = 23;
-   |         ^^^^^ the trait `Borrow<&&&&B>` is not implemented for `&B`
+   |     ^^^^^^^^^^ the trait `Borrow<&&&&B>` is not implemented for `&B`
    |
    = note: required for `HashMap<&B, u32>` to implement `Index<&&&&&B>`
 
@@ -44,10 +44,10 @@ note: required by a bound in `HashMap::<K, V, S, A>::get_mut`
   --> $SRC_DIR/std/src/collections/hash/map.rs:LL:COL
 
 error[E0277]: the trait bound `D: Borrow<&&&&D>` is not satisfied
-  --> $DIR/index-mut-help2.rs:131:9
+  --> $DIR/index-mut-help2.rs:131:5
    |
 LL |     map[index] = 23;
-   |         ^^^^^ unsatisfied trait bound
+   |     ^^^^^^^^^^ unsatisfied trait bound
    |
 help: the trait `Borrow<&&&&D>` is not implemented for `D`
       but trait `Borrow<C>` is implemented for it

tests/ui/closures/unsized-return-suggest-ref-issue-152064.fixed

@@ -7,14 +7,12 @@ fn main() {
     for s in o.map(|s| &s[3..8]) {}
     //~^ ERROR the size for values of type `str` cannot be known at compilation time
     //~| ERROR the size for values of type `str` cannot be known at compilation time
-    //~| ERROR the size for values of type `str` cannot be known at compilation time
     //~| ERROR `Option<str>` is not an iterator
 
     // Byte slice case
     let arr = Some(b"Hello, world!");
     for s in arr.map(|s| &s[3..8]) {}
     //~^ ERROR the size for values of type `[u8]` cannot be known at compilation time
     //~| ERROR the size for values of type `[u8]` cannot be known at compilation time
-    //~| ERROR the size for values of type `[u8]` cannot be known at compilation time
     //~| ERROR `Option<[u8]>` is not an iterator
 }

tests/ui/closures/unsized-return-suggest-ref-issue-152064.rs

@@ -7,14 +7,12 @@ fn main() {
     for s in o.map(|s| s[3..8]) {}
     //~^ ERROR the size for values of type `str` cannot be known at compilation time
     //~| ERROR the size for values of type `str` cannot be known at compilation time
-    //~| ERROR the size for values of type `str` cannot be known at compilation time
     //~| ERROR `Option<str>` is not an iterator
 
     // Byte slice case
     let arr = Some(b"Hello, world!");
     for s in arr.map(|s| s[3..8]) {}
     //~^ ERROR the size for values of type `[u8]` cannot be known at compilation time
     //~| ERROR the size for values of type `[u8]` cannot be known at compilation time
-    //~| ERROR the size for values of type `[u8]` cannot be known at compilation time
     //~| ERROR `Option<[u8]>` is not an iterator
 }