Auto merge of #145951 - lcnr:proof-tree-as-query, r=compiler-errors

cleanup and cache proof tree building

There's some cruft left over from when we had deep proof trees. We never encounter overflow when evaluating proof trees. Even if the recursion limit is `0`, we still only hit the overflow limit when evaluating nested goals of the root. The root goal simply inherits the `root_depth` of the `SearchGraph`.

Split `evaluate_root_goal_for_proof_tree` from the rest of the trait solver. This enables us to simplify the implementation of `evaluate_goal_raw` and the `ProofTreeBuilder` as we no longer need to manually track the state of the builder and can instead use separate types for that. It does require making a few internal methods into associated functions taking a `delegate` and a `span` instead of the `EvalCtxt` itself.

I've also split `SearchGraph::evaluate_goal` and `SearchGraph::evaluate_root_goal_for_proof_tree` for the same reason. Both functions don't actually share too much code, so by splitting them each version gets significantly easier to read.

Add a `query evaluate_root_goal_for_proof_tree_raw` to cache proof tree building. This requires arena allocating `inspect::Probe`. I've added a new type alias `I::ProbeRef` for this. We may need to adapt this for rust-analyzer? It would definitely be easy to remove the `Copy` bound here :thinking:

5 files changed, 235 insertions, 387 deletions

diff --git a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
index 4644b145b18..6f9f4067384 100644
--- a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/canonical.rs
@@ -56,22 +56,23 @@ where
     ///
     /// This expects `goal` and `opaque_types` to be eager resolved.
     pub(super) fn canonicalize_goal(
-        &self,
+        delegate: &D,
         goal: Goal<I, I::Predicate>,
         opaque_types: Vec<(ty::OpaqueTypeKey<I>, I::Ty)>,
     ) -> (Vec<I::GenericArg>, CanonicalInput<I, I::Predicate>) {
         let mut orig_values = Default::default();
         let canonical = Canonicalizer::canonicalize_input(
-            self.delegate,
+            delegate,
             &mut orig_values,
             QueryInput {
                 goal,
-                predefined_opaques_in_body: self
+                predefined_opaques_in_body: delegate
                     .cx()
                     .mk_predefined_opaques_in_body(PredefinedOpaquesData { opaque_types }),
             },
         );
-        let query_input = ty::CanonicalQueryInput { canonical, typing_mode: self.typing_mode() };
+        let query_input =
+            ty::CanonicalQueryInput { canonical, typing_mode: delegate.typing_mode() };
         (orig_values, query_input)
     }
 
@@ -271,28 +272,23 @@ where
     /// - we apply the `external_constraints` returned by the query, returning
     ///   the `normalization_nested_goals`
     pub(super) fn instantiate_and_apply_query_response(
-        &mut self,
+        delegate: &D,
         param_env: I::ParamEnv,
         original_values: &[I::GenericArg],
         response: CanonicalResponse<I>,
+        span: I::Span,
     ) -> (NestedNormalizationGoals<I>, Certainty) {
         let instantiation = Self::compute_query_response_instantiation_values(
-            self.delegate,
+            delegate,
             &original_values,
             &response,
-            self.origin_span,
+            span,
         );
 
         let Response { var_values, external_constraints, certainty } =
-            self.delegate.instantiate_canonical(response, instantiation);
+            delegate.instantiate_canonical(response, instantiation);
 
-        Self::unify_query_var_values(
-            self.delegate,
-            param_env,
-            &original_values,
-            var_values,
-            self.origin_span,
-        );
+        Self::unify_query_var_values(delegate, param_env, &original_values, var_values, span);
 
         let ExternalConstraintsData {
             region_constraints,
@@ -300,8 +296,8 @@ where
             normalization_nested_goals,
         } = &*external_constraints;
 
-        self.register_region_constraints(region_constraints);
-        self.register_new_opaque_types(opaque_types);
+        Self::register_region_constraints(delegate, region_constraints, span);
+        Self::register_new_opaque_types(delegate, opaque_types, span);
 
         (normalization_nested_goals.clone(), certainty)
     }
@@ -424,21 +420,26 @@ where
     }
 
     fn register_region_constraints(
-        &mut self,
+        delegate: &D,
         outlives: &[ty::OutlivesPredicate<I, I::GenericArg>],
+        span: I::Span,
     ) {
         for &ty::OutlivesPredicate(lhs, rhs) in outlives {
             match lhs.kind() {
-                ty::GenericArgKind::Lifetime(lhs) => self.register_region_outlives(lhs, rhs),
-                ty::GenericArgKind::Type(lhs) => self.register_ty_outlives(lhs, rhs),
+                ty::GenericArgKind::Lifetime(lhs) => delegate.sub_regions(rhs, lhs, span),
+                ty::GenericArgKind::Type(lhs) => delegate.register_ty_outlives(lhs, rhs, span),
                 ty::GenericArgKind::Const(_) => panic!("const outlives: {lhs:?}: {rhs:?}"),
             }
         }
     }
 
-    fn register_new_opaque_types(&mut self, opaque_types: &[(ty::OpaqueTypeKey<I>, I::Ty)]) {
+    fn register_new_opaque_types(
+        delegate: &D,
+        opaque_types: &[(ty::OpaqueTypeKey<I>, I::Ty)],
+        span: I::Span,
+    ) {
         for &(key, ty) in opaque_types {
-            let prev = self.delegate.register_hidden_type_in_storage(key, ty, self.origin_span);
+            let prev = delegate.register_hidden_type_in_storage(key, ty, span);
             // We eagerly resolve inference variables when computing the query response.
             // This can cause previously distinct opaque type keys to now be structurally equal.
             //
@@ -447,7 +448,7 @@ where
             // types here. However, doing so is difficult as it may result in nested goals and
             // any errors may make it harder to track the control flow for diagnostics.
             if let Some(prev) = prev {
-                self.delegate.add_duplicate_opaque_type(key, prev, self.origin_span);
+                delegate.add_duplicate_opaque_type(key, prev, span);
             }
         }
     }
diff --git a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
index 4f87902e46e..31106a74527 100644
--- a/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/eval_ctxt/mod.rs
@@ -20,13 +20,12 @@ use crate::coherence;
 use crate::delegate::SolverDelegate;
 use crate::placeholder::BoundVarReplacer;
 use crate::resolve::eager_resolve_vars;
-use crate::solve::inspect::{self, ProofTreeBuilder};
 use crate::solve::search_graph::SearchGraph;
 use crate::solve::ty::may_use_unstable_feature;
 use crate::solve::{
-    CanonicalInput, Certainty, FIXPOINT_STEP_LIMIT, Goal, GoalEvaluation, GoalEvaluationKind,
-    GoalSource, GoalStalledOn, HasChanged, NestedNormalizationGoals, NoSolution, QueryInput,
-    QueryResult,
+    CanonicalInput, Certainty, FIXPOINT_STEP_LIMIT, Goal, GoalEvaluation, GoalSource,
+    GoalStalledOn, HasChanged, NestedNormalizationGoals, NoSolution, QueryInput, QueryResult,
+    inspect,
 };
 
 pub(super) mod canonical;
@@ -130,7 +129,7 @@ where
     // evaluation code.
     tainted: Result<(), NoSolution>,
 
-    pub(super) inspect: ProofTreeBuilder<D>,
+    pub(super) inspect: inspect::EvaluationStepBuilder<D>,
 }
 
 #[derive(PartialEq, Eq, Debug, Hash, Clone, Copy)]
@@ -172,10 +171,7 @@ pub trait SolverDelegateEvalExt: SolverDelegate {
         goal: Goal<Self::Interner, <Self::Interner as Interner>::Predicate>,
         span: <Self::Interner as Interner>::Span,
     ) -> (
-        Result<
-            (NestedNormalizationGoals<Self::Interner>, GoalEvaluation<Self::Interner>),
-            NoSolution,
-        >,
+        Result<NestedNormalizationGoals<Self::Interner>, NoSolution>,
         inspect::GoalEvaluation<Self::Interner>,
     );
 }
@@ -192,14 +188,9 @@ where
         span: I::Span,
         stalled_on: Option<GoalStalledOn<I>>,
     ) -> Result<GoalEvaluation<I>, NoSolution> {
-        EvalCtxt::enter_root(
-            self,
-            self.cx().recursion_limit(),
-            GenerateProofTree::No,
-            span,
-            |ecx| ecx.evaluate_goal(GoalEvaluationKind::Root, GoalSource::Misc, goal, stalled_on),
-        )
-        .0
+        EvalCtxt::enter_root(self, self.cx().recursion_limit(), span, |ecx| {
+            ecx.evaluate_goal(GoalSource::Misc, goal, stalled_on)
+        })
     }
 
     fn root_goal_may_hold_with_depth(
@@ -208,10 +199,9 @@ where
         goal: Goal<Self::Interner, <Self::Interner as Interner>::Predicate>,
     ) -> bool {
         self.probe(|| {
-            EvalCtxt::enter_root(self, root_depth, GenerateProofTree::No, I::Span::dummy(), |ecx| {
-                ecx.evaluate_goal(GoalEvaluationKind::Root, GoalSource::Misc, goal, None)
+            EvalCtxt::enter_root(self, root_depth, I::Span::dummy(), |ecx| {
+                ecx.evaluate_goal(GoalSource::Misc, goal, None)
             })
-            .0
         })
         .is_ok()
     }
@@ -221,18 +211,8 @@ where
         &self,
         goal: Goal<I, I::Predicate>,
         span: I::Span,
-    ) -> (
-        Result<(NestedNormalizationGoals<I>, GoalEvaluation<I>), NoSolution>,
-        inspect::GoalEvaluation<I>,
-    ) {
-        let (result, proof_tree) = EvalCtxt::enter_root(
-            self,
-            self.cx().recursion_limit(),
-            GenerateProofTree::Yes,
-            span,
-            |ecx| ecx.evaluate_goal_raw(GoalEvaluationKind::Root, GoalSource::Misc, goal, None),
-        );
-        (result, proof_tree.unwrap())
+    ) -> (Result<NestedNormalizationGoals<I>, NoSolution>, inspect::GoalEvaluation<I>) {
+        evaluate_root_goal_for_proof_tree(self, goal, span)
     }
 }
 
@@ -301,17 +281,16 @@ where
     pub(super) fn enter_root<R>(
         delegate: &D,
         root_depth: usize,
-        generate_proof_tree: GenerateProofTree,
         origin_span: I::Span,
         f: impl FnOnce(&mut EvalCtxt<'_, D>) -> R,
-    ) -> (R, Option<inspect::GoalEvaluation<I>>) {
+    ) -> R {
         let mut search_graph = SearchGraph::new(root_depth);
 
         let mut ecx = EvalCtxt {
             delegate,
             search_graph: &mut search_graph,
             nested_goals: Default::default(),
-            inspect: ProofTreeBuilder::new_maybe_root(generate_proof_tree),
+            inspect: inspect::EvaluationStepBuilder::new_noop(),
 
             // Only relevant when canonicalizing the response,
             // which we don't do within this evaluation context.
@@ -324,15 +303,12 @@ where
             tainted: Ok(()),
         };
         let result = f(&mut ecx);
-
-        let proof_tree = ecx.inspect.finalize();
         assert!(
             ecx.nested_goals.is_empty(),
             "root `EvalCtxt` should not have any goals added to it"
         );
-
         assert!(search_graph.is_empty());
-        (result, proof_tree)
+        result
     }
 
     /// Creates a nested evaluation context that shares the same search graph as the
@@ -346,11 +322,10 @@ where
         cx: I,
         search_graph: &'a mut SearchGraph<D>,
         canonical_input: CanonicalInput<I>,
-        canonical_goal_evaluation: &mut ProofTreeBuilder<D>,
+        proof_tree_builder: &mut inspect::ProofTreeBuilder<D>,
         f: impl FnOnce(&mut EvalCtxt<'_, D>, Goal<I, I::Predicate>) -> R,
     ) -> R {
         let (ref delegate, input, var_values) = D::build_with_canonical(cx, &canonical_input);
-
         for &(key, ty) in &input.predefined_opaques_in_body.opaque_types {
             let prev = delegate.register_hidden_type_in_storage(key, ty, I::Span::dummy());
             // It may be possible that two entries in the opaque type storage end up
@@ -381,12 +356,12 @@ where
             nested_goals: Default::default(),
             origin_span: I::Span::dummy(),
             tainted: Ok(()),
-            inspect: canonical_goal_evaluation.new_goal_evaluation_step(var_values),
+            inspect: proof_tree_builder.new_evaluation_step(var_values),
         };
 
         let result = f(&mut ecx, input.goal);
         ecx.inspect.probe_final_state(ecx.delegate, ecx.max_input_universe);
-        canonical_goal_evaluation.goal_evaluation_step(ecx.inspect);
+        proof_tree_builder.finish_evaluation_step(ecx.inspect);
 
         // When creating a query response we clone the opaque type constraints
         // instead of taking them. This would cause an ICE here, since we have
@@ -406,13 +381,12 @@ where
     /// been constrained and the certainty of the result.
     fn evaluate_goal(
         &mut self,
-        goal_evaluation_kind: GoalEvaluationKind,
         source: GoalSource,
         goal: Goal<I, I::Predicate>,
         stalled_on: Option<GoalStalledOn<I>>,
     ) -> Result<GoalEvaluation<I>, NoSolution> {
         let (normalization_nested_goals, goal_evaluation) =
-            self.evaluate_goal_raw(goal_evaluation_kind, source, goal, stalled_on)?;
+            self.evaluate_goal_raw(source, goal, stalled_on)?;
         assert!(normalization_nested_goals.is_empty());
         Ok(goal_evaluation)
     }
@@ -426,7 +400,6 @@ where
     /// storage.
     pub(super) fn evaluate_goal_raw(
         &mut self,
-        goal_evaluation_kind: GoalEvaluationKind,
         source: GoalSource,
         goal: Goal<I, I::Predicate>,
         stalled_on: Option<GoalStalledOn<I>>,
@@ -458,17 +431,14 @@ where
         let opaque_types = self.delegate.clone_opaque_types_lookup_table();
         let (goal, opaque_types) = eager_resolve_vars(self.delegate, (goal, opaque_types));
 
-        let (orig_values, canonical_goal) = self.canonicalize_goal(goal, opaque_types);
-        let mut goal_evaluation =
-            self.inspect.new_goal_evaluation(goal, &orig_values, goal_evaluation_kind);
+        let (orig_values, canonical_goal) =
+            Self::canonicalize_goal(self.delegate, goal, opaque_types);
         let canonical_result = self.search_graph.evaluate_goal(
             self.cx(),
             canonical_goal,
             self.step_kind_for_source(source),
-            &mut goal_evaluation,
+            &mut inspect::ProofTreeBuilder::new_noop(),
         );
-        goal_evaluation.query_result(canonical_result);
-        self.inspect.goal_evaluation(goal_evaluation);
         let response = match canonical_result {
             Err(e) => return Err(e),
             Ok(response) => response,
@@ -477,8 +447,13 @@ where
         let has_changed =
             if !has_only_region_constraints(response) { HasChanged::Yes } else { HasChanged::No };
 
-        let (normalization_nested_goals, certainty) =
-            self.instantiate_and_apply_query_response(goal.param_env, &orig_values, response);
+        let (normalization_nested_goals, certainty) = Self::instantiate_and_apply_query_response(
+            self.delegate,
+            goal.param_env,
+            &orig_values,
+            response,
+            self.origin_span,
+        );
 
         // FIXME: We previously had an assert here that checked that recomputing
         // a goal after applying its constraints did not change its response.
@@ -676,12 +651,7 @@ where
                 let (
                     NestedNormalizationGoals(nested_goals),
                     GoalEvaluation { goal, certainty, stalled_on, has_changed: _ },
-                ) = self.evaluate_goal_raw(
-                    GoalEvaluationKind::Nested,
-                    source,
-                    unconstrained_goal,
-                    stalled_on,
-                )?;
+                ) = self.evaluate_goal_raw(source, unconstrained_goal, stalled_on)?;
                 // Add the nested goals from normalization to our own nested goals.
                 trace!(?nested_goals);
                 self.nested_goals.extend(nested_goals.into_iter().map(|(s, g)| (s, g, None)));
@@ -734,7 +704,7 @@ where
                 }
             } else {
                 let GoalEvaluation { goal, certainty, has_changed, stalled_on } =
-                    self.evaluate_goal(GoalEvaluationKind::Nested, source, goal, stalled_on)?;
+                    self.evaluate_goal(source, goal, stalled_on)?;
                 if has_changed == HasChanged::Yes {
                     unchanged_certainty = None;
                 }
@@ -1297,3 +1267,62 @@ where
         if predicate.allow_normalization() { predicate.super_fold_with(self) } else { predicate }
     }
 }
+
+/// Do not call this directly, use the `tcx` query instead.
+pub fn evaluate_root_goal_for_proof_tree_raw_provider<
+    D: SolverDelegate<Interner = I>,
+    I: Interner,
+>(
+    cx: I,
+    canonical_goal: CanonicalInput<I>,
+) -> (QueryResult<I>, I::ProbeRef) {
+    let mut inspect = inspect::ProofTreeBuilder::new();
+    let canonical_result = SearchGraph::<D>::evaluate_root_goal_for_proof_tree(
+        cx,
+        cx.recursion_limit(),
+        canonical_goal,
+        &mut inspect,
+    );
+    let final_revision = inspect.unwrap();
+    (canonical_result, cx.mk_probe_ref(final_revision))
+}
+
+/// Evaluate a goal to build a proof tree.
+///
+/// This is a copy of [EvalCtxt::evaluate_goal_raw] which avoids relying on the
+/// [EvalCtxt] and uses a separate cache.
+pub(super) fn evaluate_root_goal_for_proof_tree<D: SolverDelegate<Interner = I>, I: Interner>(
+    delegate: &D,
+    goal: Goal<I, I::Predicate>,
+    origin_span: I::Span,
+) -> (Result<NestedNormalizationGoals<I>, NoSolution>, inspect::GoalEvaluation<I>) {
+    let opaque_types = delegate.clone_opaque_types_lookup_table();
+    let (goal, opaque_types) = eager_resolve_vars(delegate, (goal, opaque_types));
+
+    let (orig_values, canonical_goal) = EvalCtxt::canonicalize_goal(delegate, goal, opaque_types);
+
+    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,
+        final_revision,
+        result: canonical_result,
+    };
+
+    let response = match canonical_result {
+        Err(e) => return (Err(e), proof_tree),
+        Ok(response) => response,
+    };
+
+    let (normalization_nested_goals, _certainty) = EvalCtxt::instantiate_and_apply_query_response(
+        delegate,
+        goal.param_env,
+        &proof_tree.orig_values,
+        response,
+        origin_span,
+    );
+
+    (Ok(normalization_nested_goals), proof_tree)
+}
diff --git a/compiler/rustc_next_trait_solver/src/solve/inspect/build.rs b/compiler/rustc_next_trait_solver/src/solve/inspect/build.rs
index fc56b006d94..2675ed0d0da 100644
--- a/compiler/rustc_next_trait_solver/src/solve/inspect/build.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/inspect/build.rs
@@ -12,96 +12,86 @@ use rustc_type_ir::{self as ty, Interner};
 
 use crate::delegate::SolverDelegate;
 use crate::solve::eval_ctxt::canonical;
-use crate::solve::{
-    Certainty, GenerateProofTree, Goal, GoalEvaluationKind, GoalSource, QueryResult, inspect,
-};
+use crate::solve::{Certainty, Goal, GoalSource, QueryResult, inspect};
 
-/// The core data structure when building proof trees.
+/// We need to know whether to build a prove tree while evaluating. We
+/// pass a `ProofTreeBuilder` with `state: Some(None)` into the search
+/// graph which then causes the initial `EvalCtxt::compute_goal` to actually
+/// build a proof tree which then gets written into the `state`.
 ///
-/// In case the current evaluation does not generate a proof
-/// tree, `state` is simply `None` and we avoid any work.
-///
-/// The possible states of the solver are represented via
-/// variants of [DebugSolver]. For any nested computation we call
-/// `ProofTreeBuilder::new_nested_computation_kind` which
-/// creates a new `ProofTreeBuilder` to temporarily replace the
-/// current one. Once that nested computation is done,
-/// `ProofTreeBuilder::nested_computation_kind` is called
-/// to add the finished nested evaluation to the parent.
-///
-/// We provide additional information to the current state
-/// by calling methods such as `ProofTreeBuilder::probe_kind`.
-///
-/// The actual structure closely mirrors the finished proof
-/// trees. At the end of trait solving `ProofTreeBuilder::finalize`
-/// is called to recursively convert the whole structure to a
-/// finished proof tree.
+/// Building the proof tree for a single evaluation step happens via the
+/// [EvaluationStepBuilder] which is updated by the `EvalCtxt` when
+/// appropriate.
 pub(crate) struct ProofTreeBuilder<D, I = <D as SolverDelegate>::Interner>
 where
     D: SolverDelegate<Interner = I>,
     I: Interner,
 {
+    state: Option<Box<Option<inspect::Probe<I>>>>,
     _infcx: PhantomData<D>,
-    state: Option<Box<DebugSolver<I>>>,
 }
 
-/// The current state of the proof tree builder, at most places
-/// in the code, only one or two variants are actually possible.
-///
-/// We simply ICE in case that assumption is broken.
-#[derive_where(Debug; I: Interner)]
-enum DebugSolver<I: Interner> {
-    Root,
-    GoalEvaluation(WipGoalEvaluation<I>),
-    CanonicalGoalEvaluationStep(WipCanonicalGoalEvaluationStep<I>),
-}
+impl<D: SolverDelegate<Interner = I>, I: Interner> ProofTreeBuilder<D> {
+    pub(crate) fn new() -> ProofTreeBuilder<D> {
+        ProofTreeBuilder { state: Some(Box::new(None)), _infcx: PhantomData }
+    }
 
-impl<I: Interner> From<WipGoalEvaluation<I>> for DebugSolver<I> {
-    fn from(g: WipGoalEvaluation<I>) -> DebugSolver<I> {
-        DebugSolver::GoalEvaluation(g)
+    pub(crate) fn new_noop() -> ProofTreeBuilder<D> {
+        ProofTreeBuilder { state: None, _infcx: PhantomData }
     }
-}
 
-impl<I: Interner> From<WipCanonicalGoalEvaluationStep<I>> for DebugSolver<I> {
-    fn from(g: WipCanonicalGoalEvaluationStep<I>) -> DebugSolver<I> {
-        DebugSolver::CanonicalGoalEvaluationStep(g)
+    pub(crate) fn is_noop(&self) -> bool {
+        self.state.is_none()
     }
-}
 
-#[derive_where(PartialEq, Debug; I: Interner)]
-struct WipGoalEvaluation<I: Interner> {
-    pub uncanonicalized_goal: Goal<I, I::Predicate>,
-    pub orig_values: Vec<I::GenericArg>,
-    pub encountered_overflow: bool,
-    /// After we finished evaluating this is moved into `kind`.
-    pub final_revision: Option<WipCanonicalGoalEvaluationStep<I>>,
-    pub result: Option<QueryResult<I>>,
-}
+    pub(crate) fn new_evaluation_step(
+        &mut self,
+        var_values: ty::CanonicalVarValues<I>,
+    ) -> EvaluationStepBuilder<D> {
+        if self.is_noop() {
+            EvaluationStepBuilder { state: None, _infcx: PhantomData }
+        } else {
+            EvaluationStepBuilder {
+                state: Some(Box::new(WipEvaluationStep {
+                    var_values: var_values.var_values.to_vec(),
+                    evaluation: WipProbe {
+                        initial_num_var_values: var_values.len(),
+                        steps: vec![],
+                        kind: None,
+                        final_state: None,
+                    },
+                    probe_depth: 0,
+                })),
+                _infcx: PhantomData,
+            }
+        }
+    }
 
-impl<I: Interner> Eq for WipGoalEvaluation<I> {}
-
-impl<I: Interner> WipGoalEvaluation<I> {
-    fn finalize(self) -> inspect::GoalEvaluation<I> {
-        inspect::GoalEvaluation {
-            uncanonicalized_goal: self.uncanonicalized_goal,
-            orig_values: self.orig_values,
-            kind: if self.encountered_overflow {
-                assert!(self.final_revision.is_none());
-                inspect::GoalEvaluationKind::Overflow
-            } else {
-                let final_revision = self.final_revision.unwrap().finalize();
-                inspect::GoalEvaluationKind::Evaluation { final_revision }
-            },
-            result: self.result.unwrap(),
+    pub(crate) fn finish_evaluation_step(
+        &mut self,
+        goal_evaluation_step: EvaluationStepBuilder<D>,
+    ) {
+        if let Some(this) = self.state.as_deref_mut() {
+            *this = Some(goal_evaluation_step.state.unwrap().finalize());
         }
     }
+
+    pub(crate) fn unwrap(self) -> inspect::Probe<I> {
+        self.state.unwrap().unwrap()
+    }
 }
 
-/// This only exists during proof tree building and does not have
-/// a corresponding struct in `inspect`. We need this to track a
-/// bunch of metadata about the current evaluation.
-#[derive_where(PartialEq, Debug; I: Interner)]
-struct WipCanonicalGoalEvaluationStep<I: Interner> {
+pub(crate) struct EvaluationStepBuilder<D, I = <D as SolverDelegate>::Interner>
+where
+    D: SolverDelegate<Interner = I>,
+    I: Interner,
+{
+    state: Option<Box<WipEvaluationStep<I>>>,
+    _infcx: PhantomData<D>,
+}
+
+#[derive_where(PartialEq, Eq, Debug; I: Interner)]
+struct WipEvaluationStep<I: Interner> {
     /// Unlike `EvalCtxt::var_values`, we append a new
     /// generic arg here whenever we create a new inference
     /// variable.
@@ -113,9 +103,7 @@ struct WipCanonicalGoalEvaluationStep<I: Interner> {
     evaluation: WipProbe<I>,
 }
 
-impl<I: Interner> Eq for WipCanonicalGoalEvaluationStep<I> {}
-
-impl<I: Interner> WipCanonicalGoalEvaluationStep<I> {
+impl<I: Interner> WipEvaluationStep<I> {
     fn current_evaluation_scope(&mut self) -> &mut WipProbe<I> {
         let mut current = &mut self.evaluation;
         for _ in 0..self.probe_depth {
@@ -181,169 +169,48 @@ impl<I: Interner> WipProbeStep<I> {
     }
 }
 
-impl<D: SolverDelegate<Interner = I>, I: Interner> ProofTreeBuilder<D> {
-    fn new(state: impl Into<DebugSolver<I>>) -> ProofTreeBuilder<D> {
-        ProofTreeBuilder { state: Some(Box::new(state.into())), _infcx: PhantomData }
-    }
-
-    fn opt_nested<T: Into<DebugSolver<I>>>(&self, state: impl FnOnce() -> Option<T>) -> Self {
-        ProofTreeBuilder {
-            state: self.state.as_ref().and_then(|_| Some(state()?.into())).map(Box::new),
-            _infcx: PhantomData,
-        }
+impl<D: SolverDelegate<Interner = I>, I: Interner> EvaluationStepBuilder<D> {
+    pub(crate) fn new_noop() -> EvaluationStepBuilder<D> {
+        EvaluationStepBuilder { state: None, _infcx: PhantomData }
     }
 
-    fn nested<T: Into<DebugSolver<I>>>(&self, state: impl FnOnce() -> T) -> Self {
-        ProofTreeBuilder {
-            state: self.state.as_ref().map(|_| Box::new(state().into())),
-            _infcx: PhantomData,
-        }
+    pub(crate) fn is_noop(&self) -> bool {
+        self.state.is_none()
     }
 
-    fn as_mut(&mut self) -> Option<&mut DebugSolver<I>> {
+    fn as_mut(&mut self) -> Option<&mut WipEvaluationStep<I>> {
         self.state.as_deref_mut()
     }
 
-    pub(crate) fn take_and_enter_probe(&mut self) -> ProofTreeBuilder<D> {
-        let mut nested = ProofTreeBuilder { state: self.state.take(), _infcx: PhantomData };
+    pub(crate) fn take_and_enter_probe(&mut self) -> EvaluationStepBuilder<D> {
+        let mut nested = EvaluationStepBuilder { state: self.state.take(), _infcx: PhantomData };
         nested.enter_probe();
         nested
     }
 
-    pub(crate) fn finalize(self) -> Option<inspect::GoalEvaluation<I>> {
-        match *self.state? {
-            DebugSolver::GoalEvaluation(wip_goal_evaluation) => {
-                Some(wip_goal_evaluation.finalize())
-            }
-            root => unreachable!("unexpected proof tree builder root node: {:?}", root),
-        }
-    }
-
-    pub(crate) fn new_maybe_root(generate_proof_tree: GenerateProofTree) -> ProofTreeBuilder<D> {
-        match generate_proof_tree {
-            GenerateProofTree::No => ProofTreeBuilder::new_noop(),
-            GenerateProofTree::Yes => ProofTreeBuilder::new_root(),
-        }
-    }
-
-    fn new_root() -> ProofTreeBuilder<D> {
-        ProofTreeBuilder::new(DebugSolver::Root)
-    }
-
-    fn new_noop() -> ProofTreeBuilder<D> {
-        ProofTreeBuilder { state: None, _infcx: PhantomData }
-    }
-
-    pub(crate) fn is_noop(&self) -> bool {
-        self.state.is_none()
-    }
-
-    pub(in crate::solve) fn new_goal_evaluation(
-        &mut self,
-        uncanonicalized_goal: Goal<I, I::Predicate>,
-        orig_values: &[I::GenericArg],
-        kind: GoalEvaluationKind,
-    ) -> ProofTreeBuilder<D> {
-        self.opt_nested(|| match kind {
-            GoalEvaluationKind::Root => Some(WipGoalEvaluation {
-                uncanonicalized_goal,
-                orig_values: orig_values.to_vec(),
-                encountered_overflow: false,
-                final_revision: None,
-                result: None,
-            }),
-            GoalEvaluationKind::Nested => None,
-        })
-    }
-
-    pub(crate) fn canonical_goal_evaluation_overflow(&mut self) {
+    pub(crate) fn add_var_value<T: Into<I::GenericArg>>(&mut self, arg: T) {
         if let Some(this) = self.as_mut() {
-            match this {
-                DebugSolver::GoalEvaluation(goal_evaluation) => {
-                    goal_evaluation.encountered_overflow = true;
-                }
-                _ => unreachable!(),
-            };
+            this.var_values.push(arg.into());
         }
     }
 
-    pub(crate) fn goal_evaluation(&mut self, goal_evaluation: ProofTreeBuilder<D>) {
+    fn enter_probe(&mut self) {
         if let Some(this) = self.as_mut() {
-            match this {
-                DebugSolver::Root => *this = *goal_evaluation.state.unwrap(),
-                DebugSolver::CanonicalGoalEvaluationStep(_) => {
-                    assert!(goal_evaluation.state.is_none())
-                }
-                _ => unreachable!(),
-            }
-        }
-    }
-
-    pub(crate) fn new_goal_evaluation_step(
-        &mut self,
-        var_values: ty::CanonicalVarValues<I>,
-    ) -> ProofTreeBuilder<D> {
-        self.nested(|| WipCanonicalGoalEvaluationStep {
-            var_values: var_values.var_values.to_vec(),
-            evaluation: WipProbe {
-                initial_num_var_values: var_values.len(),
+            let initial_num_var_values = this.var_values.len();
+            this.current_evaluation_scope().steps.push(WipProbeStep::NestedProbe(WipProbe {
+                initial_num_var_values,
                 steps: vec![],
                 kind: None,
                 final_state: None,
-            },
-            probe_depth: 0,
-        })
-    }
-
-    pub(crate) fn goal_evaluation_step(&mut self, goal_evaluation_step: ProofTreeBuilder<D>) {
-        if let Some(this) = self.as_mut() {
-            match (this, *goal_evaluation_step.state.unwrap()) {
-                (
-                    DebugSolver::GoalEvaluation(goal_evaluation),
-                    DebugSolver::CanonicalGoalEvaluationStep(goal_evaluation_step),
-                ) => {
-                    goal_evaluation.final_revision = Some(goal_evaluation_step);
-                }
-                _ => unreachable!(),
-            }
-        }
-    }
-
-    pub(crate) fn add_var_value<T: Into<I::GenericArg>>(&mut self, arg: T) {
-        match self.as_mut() {
-            None => {}
-            Some(DebugSolver::CanonicalGoalEvaluationStep(state)) => {
-                state.var_values.push(arg.into());
-            }
-            Some(s) => panic!("tried to add var values to {s:?}"),
-        }
-    }
-
-    fn enter_probe(&mut self) {
-        match self.as_mut() {
-            None => {}
-            Some(DebugSolver::CanonicalGoalEvaluationStep(state)) => {
-                let initial_num_var_values = state.var_values.len();
-                state.current_evaluation_scope().steps.push(WipProbeStep::NestedProbe(WipProbe {
-                    initial_num_var_values,
-                    steps: vec![],
-                    kind: None,
-                    final_state: None,
-                }));
-                state.probe_depth += 1;
-            }
-            Some(s) => panic!("tried to start probe to {s:?}"),
+            }));
+            this.probe_depth += 1;
         }
     }
 
     pub(crate) fn probe_kind(&mut self, probe_kind: inspect::ProbeKind<I>) {
-        match self.as_mut() {
-            None => {}
-            Some(DebugSolver::CanonicalGoalEvaluationStep(state)) => {
-                let prev = state.current_evaluation_scope().kind.replace(probe_kind);
-                assert_eq!(prev, None);
-            }
-            _ => panic!(),
+        if let Some(this) = self.as_mut() {
+            let prev = this.current_evaluation_scope().kind.replace(probe_kind);
+            assert_eq!(prev, None);
         }
     }
 
@@ -352,19 +219,11 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> ProofTreeBuilder<D> {
         delegate: &D,
         max_input_universe: ty::UniverseIndex,
     ) {
-        match self.as_mut() {
-            None => {}
-            Some(DebugSolver::CanonicalGoalEvaluationStep(state)) => {
-                let final_state = canonical::make_canonical_state(
-                    delegate,
-                    &state.var_values,
-                    max_input_universe,
-                    (),
-                );
-                let prev = state.current_evaluation_scope().final_state.replace(final_state);
-                assert_eq!(prev, None);
-            }
-            _ => panic!(),
+        if let Some(this) = self.as_mut() {
+            let final_state =
+                canonical::make_canonical_state(delegate, &this.var_values, max_input_universe, ());
+            let prev = this.current_evaluation_scope().final_state.replace(final_state);
+            assert_eq!(prev, None);
         }
     }
 
@@ -375,18 +234,14 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> ProofTreeBuilder<D> {
         source: GoalSource,
         goal: Goal<I, I::Predicate>,
     ) {
-        match self.as_mut() {
-            None => {}
-            Some(DebugSolver::CanonicalGoalEvaluationStep(state)) => {
-                let goal = canonical::make_canonical_state(
-                    delegate,
-                    &state.var_values,
-                    max_input_universe,
-                    goal,
-                );
-                state.current_evaluation_scope().steps.push(WipProbeStep::AddGoal(source, goal))
-            }
-            _ => panic!(),
+        if let Some(this) = self.as_mut() {
+            let goal = canonical::make_canonical_state(
+                delegate,
+                &this.var_values,
+                max_input_universe,
+                goal,
+            );
+            this.current_evaluation_scope().steps.push(WipProbeStep::AddGoal(source, goal))
         }
     }
 
@@ -396,47 +251,31 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> ProofTreeBuilder<D> {
         max_input_universe: ty::UniverseIndex,
         impl_args: I::GenericArgs,
     ) {
-        match self.as_mut() {
-            Some(DebugSolver::CanonicalGoalEvaluationStep(state)) => {
-                let impl_args = canonical::make_canonical_state(
-                    delegate,
-                    &state.var_values,
-                    max_input_universe,
-                    impl_args,
-                );
-                state
-                    .current_evaluation_scope()
-                    .steps
-                    .push(WipProbeStep::RecordImplArgs { impl_args });
-            }
-            None => {}
-            _ => panic!(),
+        if let Some(this) = self.as_mut() {
+            let impl_args = canonical::make_canonical_state(
+                delegate,
+                &this.var_values,
+                max_input_universe,
+                impl_args,
+            );
+            this.current_evaluation_scope().steps.push(WipProbeStep::RecordImplArgs { impl_args });
         }
     }
 
     pub(crate) fn make_canonical_response(&mut self, shallow_certainty: Certainty) {
-        match self.as_mut() {
-            Some(DebugSolver::CanonicalGoalEvaluationStep(state)) => {
-                state
-                    .current_evaluation_scope()
-                    .steps
-                    .push(WipProbeStep::MakeCanonicalResponse { shallow_certainty });
-            }
-            None => {}
-            _ => panic!(),
+        if let Some(this) = self.as_mut() {
+            this.current_evaluation_scope()
+                .steps
+                .push(WipProbeStep::MakeCanonicalResponse { shallow_certainty });
         }
     }
 
-    pub(crate) fn finish_probe(mut self) -> ProofTreeBuilder<D> {
-        match self.as_mut() {
-            None => {}
-            Some(DebugSolver::CanonicalGoalEvaluationStep(state)) => {
-                assert_ne!(state.probe_depth, 0);
-                let num_var_values = state.current_evaluation_scope().initial_num_var_values;
-                state.var_values.truncate(num_var_values);
-                state.probe_depth -= 1;
-            }
-            _ => panic!(),
+    pub(crate) fn finish_probe(mut self) -> EvaluationStepBuilder<D> {
+        if let Some(this) = self.as_mut() {
+            assert_ne!(this.probe_depth, 0);
+            let num_var_values = this.current_evaluation_scope().initial_num_var_values;
+            this.var_values.truncate(num_var_values);
+            this.probe_depth -= 1;
         }
 
         self
@@ -444,21 +283,7 @@ impl<D: SolverDelegate<Interner = I>, I: Interner> ProofTreeBuilder<D> {
 
     pub(crate) fn query_result(&mut self, result: QueryResult<I>) {
         if let Some(this) = self.as_mut() {
-            match this {
-                DebugSolver::GoalEvaluation(goal_evaluation) => {
-                    assert_eq!(goal_evaluation.result.replace(result), None);
-                }
-                DebugSolver::CanonicalGoalEvaluationStep(evaluation_step) => {
-                    assert_eq!(
-                        evaluation_step
-                            .evaluation
-                            .kind
-                            .replace(inspect::ProbeKind::Root { result }),
-                        None
-                    );
-                }
-                _ => unreachable!(),
-            }
+            assert_eq!(this.evaluation.kind.replace(inspect::ProbeKind::Root { result }), None);
         }
     }
 }
diff --git a/compiler/rustc_next_trait_solver/src/solve/mod.rs b/compiler/rustc_next_trait_solver/src/solve/mod.rs
index 710b59f662a..85f9d852d95 100644
--- a/compiler/rustc_next_trait_solver/src/solve/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/mod.rs
@@ -27,7 +27,10 @@ pub use rustc_type_ir::solve::*;
 use rustc_type_ir::{self as ty, Interner, TypingMode};
 use tracing::instrument;
 
-pub use self::eval_ctxt::{EvalCtxt, GenerateProofTree, SolverDelegateEvalExt};
+pub use self::eval_ctxt::{
+    EvalCtxt, GenerateProofTree, SolverDelegateEvalExt,
+    evaluate_root_goal_for_proof_tree_raw_provider,
+};
 use crate::delegate::SolverDelegate;
 use crate::solve::assembly::Candidate;
 
@@ -42,12 +45,6 @@ use crate::solve::assembly::Candidate;
 /// recursion limit again. However, this feels very unlikely.
 const FIXPOINT_STEP_LIMIT: usize = 8;
 
-#[derive(Debug, Copy, Clone, PartialEq, Eq)]
-enum GoalEvaluationKind {
-    Root,
-    Nested,
-}
-
 /// Whether evaluating this goal ended up changing the
 /// inference state.
 #[derive(PartialEq, Eq, Debug, Hash, Clone, Copy)]
diff --git a/compiler/rustc_next_trait_solver/src/solve/search_graph.rs b/compiler/rustc_next_trait_solver/src/solve/search_graph.rs
index 84f8eda4f8d..f0342e0523f 100644
--- a/compiler/rustc_next_trait_solver/src/solve/search_graph.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/search_graph.rs
@@ -7,8 +7,9 @@ use rustc_type_ir::solve::{CanonicalInput, Certainty, NoSolution, QueryResult};
 use rustc_type_ir::{Interner, TypingMode};
 
 use crate::delegate::SolverDelegate;
-use crate::solve::inspect::ProofTreeBuilder;
-use crate::solve::{EvalCtxt, FIXPOINT_STEP_LIMIT, has_no_inference_or_external_constraints};
+use crate::solve::{
+    EvalCtxt, FIXPOINT_STEP_LIMIT, has_no_inference_or_external_constraints, inspect,
+};
 
 /// This type is never constructed. We only use it to implement `search_graph::Delegate`
 /// for all types which impl `SolverDelegate` and doing it directly fails in coherence.
@@ -34,7 +35,7 @@ where
 
     const FIXPOINT_STEP_LIMIT: usize = FIXPOINT_STEP_LIMIT;
 
-    type ProofTreeBuilder = ProofTreeBuilder<D>;
+    type ProofTreeBuilder = inspect::ProofTreeBuilder<D>;
     fn inspect_is_noop(inspect: &mut Self::ProofTreeBuilder) -> bool {
         inspect.is_noop()
     }
@@ -81,12 +82,7 @@ where
         Self::initial_provisional_result(cx, kind, input) == result
     }
 
-    fn on_stack_overflow(
-        cx: I,
-        input: CanonicalInput<I>,
-        inspect: &mut ProofTreeBuilder<D>,
-    ) -> QueryResult<I> {
-        inspect.canonical_goal_evaluation_overflow();
+    fn on_stack_overflow(cx: I, input: CanonicalInput<I>) -> QueryResult<I> {
         response_no_constraints(cx, input, Certainty::overflow(true))
     }