Reuse one base universe for proof tree replay

Store `max_input_universe` in `GoalEvaluation`/`InspectGoal` and
pass it to `instantiate_canonical_state` for all replay of a
single goal.

This keeps every canonical state from that goal using the same
universe mapping instead of rebuilding it from `orig_values`.

Inline universe creation for query response instantiation.

Author: VicenteGusmao

compiler/rustc_next_trait_solver/src/canonical/mod.rs

@@ -10,15 +10,14 @@
 //! [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::inherent::*;
 use rustc_type_ir::relate::solver_relating::RelateExt;
 use rustc_type_ir::{
     self as ty, Canonical, CanonicalVarKind, CanonicalVarValues, InferCtxtLike, Interner,
-    TypeFoldable, TypeSuperVisitable, TypeVisitable, TypeVisitor, TypingMode, TypingModeEqWrapper,
+    TypeFoldable, TypingMode, TypingModeEqWrapper,
 };
 use tracing::instrument;
 
@@ -151,7 +150,11 @@ 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 = create_universes_for_response(delegate, response);
+    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();
+    }
 
     compute_query_response_instantiation_values_in_universe(
         delegate,
@@ -257,106 +260,6 @@ 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
-}
-
-// Recover the caller-side max input universe for replaying a canonical state.
-//
-// We want the base universe the canonical state was relative to when it was
-// first recorded, not the max universe of the current resolved contents of
-// `orig_values`. Resolved inference vars may now point at placeholders from
-// later replay work, so looking through them would overestimate the base
-// universe and shift the replayed placeholders.
-fn max_input_universe_for_canonical_state<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
-}
-
 /// 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
@@ -461,6 +364,7 @@ pub fn instantiate_canonical_state<D, I, T>(
     delegate: &D,
     span: I::Span,
     param_env: I::ParamEnv,
+    prev_universe: ty::UniverseIndex,
     orig_values: &mut Vec<I::GenericArg>,
     state: inspect::CanonicalState<I, T>,
 ) -> T
@@ -471,7 +375,6 @@ where
 {
     // 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_canonical_state(delegate, orig_values);
     let max_universe = prev_universe + state.max_universe.index();
     while delegate.universe() < max_universe {
         delegate.create_next_universe();

compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs

@@ -1789,13 +1789,15 @@ pub(super) fn evaluate_root_goal_for_proof_tree<D: SolverDelegate<Interner = I>,
 
     let (orig_values, canonical_goal) =
         canonicalize_goal(delegate, goal, &opaque_types, typing_mode.into());
+    let max_input_universe = canonical_goal.canonical.max_universe;
 
     let (canonical_result, final_revision) =
         delegate.cx().evaluate_root_goal_for_proof_tree_raw(canonical_goal);
 
     let proof_tree = inspect::GoalEvaluation {
         uncanonicalized_goal: goal,
         orig_values,
+        max_input_universe,
         final_revision,
         result: canonical_result,
     };

compiler/rustc_trait_selection/src/solve/inspect/analyse.rs

@@ -32,6 +32,7 @@ pub struct InspectGoal<'a, 'tcx> {
     infcx: &'a SolverDelegate<'tcx>,
     depth: usize,
     orig_values: Vec<ty::GenericArg<'tcx>>,
+    max_input_universe: ty::UniverseIndex,
     goal: Goal<'tcx, ty::Predicate<'tcx>>,
     result: Result<Certainty, NoSolution>,
     final_revision: &'tcx inspect::Probe<TyCtxt<'tcx>>,
@@ -102,16 +103,29 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
             match **step {
                 inspect::ProbeStep::AddGoal(source, goal) => instantiated_goals.push((
                     source,
-                    instantiate_canonical_state(infcx, span, param_env, &mut orig_values, goal),
+                    instantiate_canonical_state(
+                        infcx,
+                        span,
+                        param_env,
+                        self.goal.max_input_universe,
+                        &mut orig_values,
+                        goal,
+                    ),
                 )),
                 inspect::ProbeStep::RecordImplArgs { .. } => {}
                 inspect::ProbeStep::MakeCanonicalResponse { .. }
                 | inspect::ProbeStep::NestedProbe(_) => unreachable!(),
             }
         }
 
-        let () =
-            instantiate_canonical_state(infcx, span, param_env, &mut orig_values, self.final_state);
+        let () = instantiate_canonical_state(
+            infcx,
+            span,
+            param_env,
+            self.goal.max_input_universe,
+            &mut orig_values,
+            self.final_state,
+        );
 
         instantiated_goals
             .into_iter()
@@ -139,6 +153,7 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
                         infcx,
                         span,
                         param_env,
+                        self.goal.max_input_universe,
                         &mut orig_values,
                         impl_args,
                     );
@@ -147,6 +162,7 @@ impl<'a, 'tcx> InspectCandidate<'a, 'tcx> {
                         infcx,
                         span,
                         param_env,
+                        self.goal.max_input_universe,
                         &mut orig_values,
                         self.final_state,
                     );
@@ -321,8 +337,13 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
     ) -> Self {
         let infcx = <&SolverDelegate<'tcx>>::from(infcx);
 
-        let inspect::GoalEvaluation { uncanonicalized_goal, orig_values, final_revision, result } =
-            root;
+        let inspect::GoalEvaluation {
+            uncanonicalized_goal,
+            orig_values,
+            max_input_universe,
+            final_revision,
+            result,
+        } = root;
         // If there's a normalizes-to goal, AND the evaluation result with the result of
         // constraining the normalizes-to RHS and computing the nested goals.
         let result = result.map(|ok| ok.value.certainty);
@@ -331,6 +352,7 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
             infcx,
             depth,
             orig_values,
+            max_input_universe,
             goal: eager_resolve_vars(&**infcx, uncanonicalized_goal),
             result,
             final_revision,

compiler/rustc_type_ir/src/solve/inspect.rs

@@ -49,6 +49,7 @@ pub type CanonicalState<I, T> = Canonical<I, State<I, T>>;
 pub struct GoalEvaluation<I: Interner> {
     pub uncanonicalized_goal: Goal<I, I::Predicate>,
     pub orig_values: Vec<I::GenericArg>,
+    pub max_input_universe: crate::UniverseIndex,
     pub final_revision: I::Probe,
     pub result: QueryResult<I>,
 }