Instead of falling back in diagnostics, the fix now addresses the actual universe mismatch during proof tree replay.

Since my last implementation was not really fixing the problem I removed the functions: try_unify_query_var_values, try_instantiate_canonical_state and InspectCandidate::try_instantiate_nested_goals. I also reverted the diagnostic callers in derive_errors.rs to use instantiate_nested_goals again.

For my new implementation, in the rustc_trait_selection  I changed fresh_var_for_kind_with_span to fresh_var_for_kind_in_universe which now uses next_region_var_in_universe, next_ty_var_in_universe and next_const_var_in_universe. Also in instantiate_canonical_state it is now ensured that infcx has all the universes up to prev_universe + state.max_universe, and creates any replay-only fresh vars in that explicit universe.  I also split out compute_query_response_instantiation_values_in_universe so replay uses that same explicit base universe when reconstructing placeholders from the canonical state.

Signed-off-by: Vicente Gusmão <vicente.gusmao@tecnico.ulisboa.pt>

Author: VicenteGusmao

compiler/rustc_next_trait_solver/src/canonical/mod.rs

@@ -10,15 +10,15 @@
 //! [c]: https://rustc-dev-guide.rust-lang.org/solve/canonicalization.html
 
 use std::iter;
+use std::marker::PhantomData;
 
 use canonicalizer::Canonicalizer;
 use rustc_index::IndexVec;
-use rustc_type_ir::error::TypeError;
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::relate::solver_relating::RelateExt;
 use rustc_type_ir::{
     self as ty, Canonical, CanonicalVarKind, CanonicalVarValues, InferCtxtLike, Interner,
-    TypeFoldable, TypingMode, TypingModeEqWrapper,
+    TypeFoldable, TypeSuperVisitable, TypeVisitable, TypeVisitor, TypingMode, TypingModeEqWrapper,
 };
 use tracing::instrument;
 
@@ -145,12 +145,29 @@ where
     // FIXME: Longterm canonical queries should deal with all placeholders
     // created inside of the query directly instead of returning them to the
     // caller.
-    let prev_universe = delegate.universe();
-    let universes_created_in_query = response.max_universe.index();
-    for _ in 0..universes_created_in_query {
-        delegate.create_next_universe();
-    }
+    let prev_universe = create_universes_for_response(delegate, response);
 
+    compute_query_response_instantiation_values_in_universe(
+        delegate,
+        original_values,
+        response,
+        span,
+        prev_universe,
+    )
+}
+
+fn compute_query_response_instantiation_values_in_universe<D, I, T>(
+    delegate: &D,
+    original_values: &[I::GenericArg],
+    response: &Canonical<I, T>,
+    span: I::Span,
+    prev_universe: ty::UniverseIndex,
+) -> CanonicalVarValues<I>
+where
+    D: SolverDelegate<Interner = I>,
+    I: Interner,
+    T: ResponseT<I>,
+{
     let var_values = response.value.var_values();
     assert_eq!(original_values.len(), var_values.len());
 
@@ -234,6 +251,109 @@ where
     })
 }
 
+fn create_universes_for_response<D, I, T>(
+    delegate: &D,
+    response: &Canonical<I, T>,
+) -> ty::UniverseIndex
+where
+    D: SolverDelegate<Interner = I>,
+    I: Interner,
+{
+    let prev_universe = delegate.universe();
+    let universes_created_in_query = response.max_universe.index();
+    for _ in 0..universes_created_in_query {
+        delegate.create_next_universe();
+    }
+    prev_universe
+}
+
+fn max_input_universe_for_replay<D, I>(
+    delegate: &D,
+    orig_values: &[I::GenericArg],
+) -> ty::UniverseIndex
+where
+    D: SolverDelegate<Interner = I>,
+    I: Interner,
+{
+    struct MaxUniverseVisitor<'a, D, I> {
+        delegate: &'a D,
+        max_universe: ty::UniverseIndex,
+        _interner: PhantomData<I>,
+    }
+
+    impl<D, I> TypeVisitor<I> for MaxUniverseVisitor<'_, D, I>
+    where
+        D: SolverDelegate<Interner = I>,
+        I: Interner,
+    {
+        type Result = ();
+
+        fn visit_ty(&mut self, t: I::Ty) -> Self::Result {
+            match t.kind() {
+                ty::Placeholder(placeholder) => {
+                    self.max_universe = self.max_universe.max(placeholder.universe());
+                }
+                ty::Infer(ty::TyVar(vid)) => {
+                    if let Some(universe) = self.delegate.universe_of_ty(vid) {
+                        self.max_universe = self.max_universe.max(universe);
+                    }
+                }
+                _ => {}
+            }
+            t.super_visit_with(self)
+        }
+
+        fn visit_region(&mut self, r: I::Region) -> Self::Result {
+            match r.kind() {
+                ty::RePlaceholder(placeholder) => {
+                    self.max_universe = self.max_universe.max(placeholder.universe());
+                }
+                ty::ReVar(vid) => {
+                    if let Some(universe) = self.delegate.universe_of_lt(vid) {
+                        self.max_universe = self.max_universe.max(universe);
+                    }
+                }
+                _ => {}
+            }
+        }
+
+        fn visit_const(&mut self, c: I::Const) -> Self::Result {
+            match c.kind() {
+                ty::ConstKind::Placeholder(placeholder) => {
+                    self.max_universe = self.max_universe.max(placeholder.universe());
+                }
+                ty::ConstKind::Infer(ty::InferConst::Var(vid)) => {
+                    if let Some(universe) = self.delegate.universe_of_ct(vid) {
+                        self.max_universe = self.max_universe.max(universe);
+                    }
+                }
+                _ => {}
+            }
+            c.super_visit_with(self)
+        }
+    }
+
+    let mut visitor = MaxUniverseVisitor {
+        delegate,
+        max_universe: ty::UniverseIndex::ROOT,
+        _interner: PhantomData,
+    };
+    for value in orig_values {
+        value.visit_with(&mut visitor);
+    }
+    visitor.max_universe
+}
+
+fn ensure_universes_through<D, I>(delegate: &D, max_universe: ty::UniverseIndex)
+where
+    D: SolverDelegate<Interner = I>,
+    I: Interner,
+{
+    while delegate.universe() < max_universe {
+        delegate.create_next_universe();
+    }
+}
+
 /// Unify the `original_values` with the `var_values` returned by the canonical query..
 ///
 /// This assumes that this unification will always succeed. This is the case when
@@ -247,39 +367,23 @@ where
 /// from the solver we assume that the solver correctly handled aliases and therefore
 /// always relate them structurally here.
 #[instrument(level = "trace", skip(delegate))]
-fn try_unify_query_var_values<D, I>(
+fn unify_query_var_values<D, I>(
     delegate: &D,
     param_env: I::ParamEnv,
     original_values: &[I::GenericArg],
     var_values: CanonicalVarValues<I>,
     span: I::Span,
-) -> Result<(), TypeError<I>>
-where
+) where
     D: SolverDelegate<Interner = I>,
     I: Interner,
 {
     assert_eq!(original_values.len(), var_values.len());
 
     for (&orig, response) in iter::zip(original_values, var_values.var_values.iter()) {
-        let goals = delegate.eq_structurally_relating_aliases(param_env, orig, response, span)?;
+        let goals =
+            delegate.eq_structurally_relating_aliases(param_env, orig, response, span).unwrap();
         assert!(goals.is_empty());
     }
-
-    Ok(())
-}
-
-#[instrument(level = "trace", skip(delegate))]
-fn unify_query_var_values<D, I>(
-    delegate: &D,
-    param_env: I::ParamEnv,
-    original_values: &[I::GenericArg],
-    var_values: CanonicalVarValues<I>,
-    span: I::Span,
-) where
-    D: SolverDelegate<Interner = I>,
-    I: Interner,
-{
-    try_unify_query_var_values(delegate, param_env, original_values, var_values, span).unwrap();
 }
 
 fn register_region_constraints<D, I>(
@@ -357,43 +461,34 @@ pub fn instantiate_canonical_state<D, I, T>(
     orig_values: &mut Vec<I::GenericArg>,
     state: inspect::CanonicalState<I, T>,
 ) -> T
-where
-    D: SolverDelegate<Interner = I>,
-    I: Interner,
-    T: TypeFoldable<I>,
-{
-    try_instantiate_canonical_state(delegate, span, param_env, orig_values, state).unwrap()
-}
-
-// FIXME: needs to be pub to be accessed by downstream
-// `rustc_trait_selection::solve::inspect::analyse`.
-pub fn try_instantiate_canonical_state<D, I, T>(
-    delegate: &D,
-    span: I::Span,
-    param_env: I::ParamEnv,
-    orig_values: &mut Vec<I::GenericArg>,
-    state: inspect::CanonicalState<I, T>,
-) -> Result<T, TypeError<I>>
 where
     D: SolverDelegate<Interner = I>,
     I: Interner,
     T: TypeFoldable<I>,
 {
     // In case any fresh inference variables have been created between `state`
     // and the previous instantiation, extend `orig_values` for it.
+    let prev_universe = max_input_universe_for_replay(delegate, orig_values);
+    let max_universe = prev_universe + state.max_universe.index();
+    ensure_universes_through(delegate, max_universe);
     orig_values.extend(
         state.value.var_values.var_values.as_slice()[orig_values.len()..]
             .iter()
-            .map(|&arg| delegate.fresh_var_for_kind_with_span(arg, span)),
+            .map(|&arg| delegate.fresh_var_for_kind_in_universe(arg, span, max_universe)),
     );
 
-    let instantiation =
-        compute_query_response_instantiation_values(delegate, orig_values, &state, span);
+    let instantiation = compute_query_response_instantiation_values_in_universe(
+        delegate,
+        orig_values,
+        &state,
+        span,
+        prev_universe,
+    );
 
     let inspect::State { var_values, data } = delegate.instantiate_canonical(state, instantiation);
 
-    try_unify_query_var_values(delegate, param_env, orig_values, var_values, span)?;
-    Ok(data)
+    unify_query_var_values(delegate, param_env, orig_values, var_values, span);
+    data
 }
 
 pub fn response_no_constraints_raw<I: Interner>(

compiler/rustc_next_trait_solver/src/delegate.rs

@@ -25,10 +25,11 @@ pub trait SolverDelegate: Deref<Target = Self::Infcx> + Sized {
         span: <Self::Interner as Interner>::Span,
     ) -> Option<Certainty>;
 
-    fn fresh_var_for_kind_with_span(
+    fn fresh_var_for_kind_in_universe(
         &self,
         arg: <Self::Interner as Interner>::GenericArg,
         span: <Self::Interner as Interner>::Span,
+        universe: ty::UniverseIndex,
     ) -> <Self::Interner as Interner>::GenericArg;
 
     // FIXME: Uplift the leak check into this crate.

compiler/rustc_trait_selection/src/solve/delegate.rs

@@ -166,17 +166,18 @@ impl<'tcx> rustc_next_trait_solver::delegate::SolverDelegate for SolverDelegate<
         }
     }
 
-    fn fresh_var_for_kind_with_span(
+    fn fresh_var_for_kind_in_universe(
         &self,
         arg: ty::GenericArg<'tcx>,
         span: Span,
+        universe: ty::UniverseIndex,
     ) -> ty::GenericArg<'tcx> {
         match arg.kind() {
             ty::GenericArgKind::Lifetime(_) => {
-                self.next_region_var(RegionVariableOrigin::Misc(span)).into()
+                self.next_region_var_in_universe(RegionVariableOrigin::Misc(span), universe).into()
             }
-            ty::GenericArgKind::Type(_) => self.next_ty_var(span).into(),
-            ty::GenericArgKind::Const(_) => self.next_const_var(span).into(),
+            ty::GenericArgKind::Type(_) => self.next_ty_var_in_universe(span, universe).into(),
+            ty::GenericArgKind::Const(_) => self.next_const_var_in_universe(span, universe).into(),
         }
     }
 

compiler/rustc_trait_selection/src/solve/fulfill/derive_errors.rs

@@ -235,20 +235,14 @@ impl<'tcx> BestObligation<'tcx> {
                 if candidates.len() > 1 {
                     candidates.retain(|candidate| {
                         goal.infcx().probe(|_| {
-                            // Replaying proof-tree state is only a diagnostics aid. If the
-                            // canonical state no longer reproduces cleanly in the current
-                            // inference context, keep the candidate rather than changing the
-                            // observable error into an ICE.
-                            candidate.try_instantiate_nested_goals(self.span()).is_none_or(
-                                |nested_goals| {
-                                    nested_goals.iter().any(|nested_goal| {
-                                        matches!(
-                                            nested_goal.source(),
-                                            GoalSource::ImplWhereBound
-                                                | GoalSource::AliasBoundConstCondition
-                                                | GoalSource::AliasWellFormed
-                                        ) && nested_goal.result().is_err()
-                                    })
+                            candidate.instantiate_nested_goals(self.span()).iter().any(
+                                |nested_goal| {
+                                    matches!(
+                                        nested_goal.source(),
+                                        GoalSource::ImplWhereBound
+                                            | GoalSource::AliasBoundConstCondition
+                                            | GoalSource::AliasWellFormed
+                                    ) && nested_goal.result().is_err()
                                 },
                             )
                         })
@@ -491,12 +485,7 @@ impl<'tcx> ProofTreeVisitor<'tcx> for BestObligation<'tcx> {
             _ => ChildMode::PassThrough,
         };
 
-        let Some(nested_goals) = candidate.try_instantiate_nested_goals(self.span()) else {
-            // We failed to replay this proof-tree candidate in the current inference state.
-            // Fall back to the current obligation instead of trying to derive a more specific
-            // nested error from diagnostics-only reconstruction.
-            return ControlFlow::Break(self.obligation.clone());
-        };
+        let nested_goals = candidate.instantiate_nested_goals(self.span());
 
         // If the candidate requires some `T: FnPtr` bound which does not hold should not be treated as
         // an actual candidate, instead we should treat them as if the impl was never considered to