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

compiler/rustc_hir_typeck/src/fn_ctxt/checks.rs

@@ -17,10 +17,13 @@ use rustc_hir::{Expr, ExprKind, FnRetTy, HirId, Node, QPath, is_range_literal};
 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::print::with_forced_trimmed_paths;
+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;
@@ -250,41 +253,66 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 // function where-bounds don't actually hold. This results
                 // in weird bugs when later treating these expectations as if
                 // they were actually correct.
-                self.fudge_inference_if_ok(|| {
-                    let ocx = ObligationCtxt::new(self);
-
-                    // Attempt to apply a subtyping relationship between the formal
-                    // return type (likely containing type variables if the function
-                    // is polymorphic) and the expected return type.
-                    // No argument expectations are produced if unification fails.
-                    let origin = self.misc(call_span);
-                    ocx.sup(&origin, self.param_env, expected_output, formal_output)?;
-
-                    // Check the well-formedness of expected input tys, as using ill-formed
-                    // expectation may cause type inference errors, see #150316.
-                    for &ty in formal_input_tys {
-                        ocx.register_obligation(traits::Obligation::new(
-                            self.tcx,
-                            self.misc(call_span),
-                            self.param_env,
-                            ty::ClauseKind::WellFormed(ty.into()),
-                        ));
-                    }
+                let expected_input_tys = self
+                    .fudge_inference_if_ok(|| {
+                        let ocx = ObligationCtxt::new(self);
+
+                        // Attempt to apply a subtyping relationship between the formal
+                        // return type (likely containing type variables if the function
+                        // is polymorphic) and the expected return type.
+                        // No argument expectations are produced if unification fails.
+                        let origin = self.misc(call_span);
+                        ocx.sup(&origin, self.param_env, expected_output, formal_output)?;
+
+                        // Check the well-formedness of expected input tys, as using ill-formed
+                        // expectation may cause type inference errors, see #150316.
+                        for &ty in formal_input_tys {
+                            ocx.register_obligation(traits::Obligation::new(
+                                self.tcx,
+                                self.misc(call_span),
+                                self.param_env,
+                                ty::ClauseKind::WellFormed(ty.into()),
+                            ));
+                        }
 
-                    if !ocx.try_evaluate_obligations().no_errors() {
-                        return Err(TypeError::Mismatch);
-                    }
+                        if !ocx.try_evaluate_obligations().no_errors() {
+                            return Err(TypeError::Mismatch);
+                        }
 
-                    // Record all the argument types, with the args
-                    // produced from the above subtyping unification.
-                    Ok(Some(
-                        formal_input_tys
-                            .iter()
-                            .map(|&ty| self.resolve_vars_if_possible(ty))
-                            .collect(),
-                    ))
-                })
-                .ok()
+                        // Record all the argument types, with the args
+                        // produced from the above subtyping unification.
+                        Ok(Some(
+                            formal_input_tys
+                                .iter()
+                                .map(|&ty| self.resolve_vars_if_possible(ty))
+                                .collect::<Vec<_>>(),
+                        ))
+                    })
+                    .ok()?;
+
+                Some(expected_input_tys.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. Usually fudging helps because it gains
+                        // information from a callsite of a function. However, Fudging also sometimes
+                        // loses information, when the original, formal, input type had constraints on it,
+                        // and fudging replaces all inference variables with fresh ones, those constraints
+                        // are discarded. This check makes sure we only keep fudging output if structural
+                        // changes were made to the type. If all that was changed were some typevars,
+                        // we go back to the unfudged formal input type.
+                        .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()
+                }))
             })
             .unwrap_or_default();
 
@@ -3535,3 +3563,100 @@ 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);
+            }
+
+            // Avoid ICEs when in gce, and `structurally_relate_consts`
+            // turns a non-infer const into an infer const
+            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

@@ -1247,22 +1247,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)
+                            }
+                        }
                     }
                 }
 
@@ -1276,13 +1299,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,
             },
 
@@ -1307,6 +1333,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);
     }

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
 }