evaluate_goal: avoid unnecessary step

Author: lcnr

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

@@ -3,7 +3,7 @@ use std::ops::ControlFlow;
 
 #[cfg(feature = "nightly")]
 use rustc_macros::HashStable_NoContext;
-use rustc_type_ir::data_structures::{HashMap, HashSet, ensure_sufficient_stack};
+use rustc_type_ir::data_structures::{HashMap, HashSet};
 use rustc_type_ir::fast_reject::DeepRejectCtxt;
 use rustc_type_ir::inherent::*;
 use rustc_type_ir::relate::Relate;
@@ -339,13 +339,12 @@ where
 
     /// Creates a nested evaluation context that shares the same search graph as the
     /// one passed in. This is suitable for evaluation, granted that the search graph
-    /// has had the nested goal recorded on its stack ([`SearchGraph::with_new_goal`]),
-    /// but it's preferable to use other methods that call this one rather than this
-    /// method directly.
+    /// has had the nested goal recorded on its stack. This method only be used by
+    /// `search_graph::Delegate::compute_goal`.
     ///
     /// This function takes care of setting up the inference context, setting the anchor,
     /// and registering opaques from the canonicalized input.
-    fn enter_canonical<R>(
+    pub(super) fn enter_canonical<R>(
         cx: I,
         search_graph: &'a mut SearchGraph<D>,
         canonical_input: CanonicalInput<I>,
@@ -401,56 +400,6 @@ where
         result
     }
 
-    /// The entry point of the solver.
-    ///
-    /// This function deals with (coinductive) cycles, overflow, and caching
-    /// and then calls [`EvalCtxt::compute_goal`] which contains the actual
-    /// logic of the solver.
-    ///
-    /// Instead of calling this function directly, use either [EvalCtxt::evaluate_goal]
-    /// if you're inside of the solver or [SolverDelegateEvalExt::evaluate_root_goal] if you're
-    /// outside of it.
-    #[instrument(level = "debug", skip(cx, search_graph, goal_evaluation), ret)]
-    fn evaluate_canonical_goal(
-        cx: I,
-        search_graph: &'a mut SearchGraph<D>,
-        canonical_input: CanonicalInput<I>,
-        step_kind_from_parent: PathKind,
-        goal_evaluation: &mut ProofTreeBuilder<D>,
-    ) -> QueryResult<I> {
-        let mut canonical_goal_evaluation =
-            goal_evaluation.new_canonical_goal_evaluation(canonical_input);
-
-        // Deal with overflow, caching, and coinduction.
-        //
-        // The actual solver logic happens in `ecx.compute_goal`.
-        let result = ensure_sufficient_stack(|| {
-            search_graph.with_new_goal(
-                cx,
-                canonical_input,
-                step_kind_from_parent,
-                &mut canonical_goal_evaluation,
-                |search_graph, cx, canonical_input, canonical_goal_evaluation| {
-                    EvalCtxt::enter_canonical(
-                        cx,
-                        search_graph,
-                        canonical_input,
-                        canonical_goal_evaluation,
-                        |ecx, goal| {
-                            let result = ecx.compute_goal(goal);
-                            ecx.inspect.query_result(result);
-                            result
-                        },
-                    )
-                },
-            )
-        });
-
-        canonical_goal_evaluation.query_result(result);
-        goal_evaluation.canonical_goal_evaluation(canonical_goal_evaluation);
-        result
-    }
-
     /// Recursively evaluates `goal`, returning whether any inference vars have
     /// been constrained and the certainty of the result.
     fn evaluate_goal(
@@ -504,13 +453,16 @@ where
         let (orig_values, canonical_goal) = self.canonicalize_goal(goal);
         let mut goal_evaluation =
             self.inspect.new_goal_evaluation(goal, &orig_values, goal_evaluation_kind);
-        let canonical_response = EvalCtxt::evaluate_canonical_goal(
+        let mut canonical_goal_evaluation =
+            goal_evaluation.new_canonical_goal_evaluation(canonical_goal);
+        let canonical_response = self.search_graph.evaluate_goal(
             self.cx(),
-            self.search_graph,
             canonical_goal,
             self.step_kind_for_source(source),
-            &mut goal_evaluation,
+            &mut canonical_goal_evaluation,
         );
+        canonical_goal_evaluation.query_result(canonical_response);
+        goal_evaluation.canonical_goal_evaluation(canonical_goal_evaluation);
         let response = match canonical_response {
             Err(e) => {
                 self.inspect.goal_evaluation(goal_evaluation);
@@ -585,7 +537,7 @@ where
         Ok((normalization_nested_goals, GoalEvaluation { certainty, has_changed, stalled_on }))
     }
 
-    fn compute_goal(&mut self, goal: Goal<I, I::Predicate>) -> QueryResult<I> {
+    pub(super) fn compute_goal(&mut self, goal: Goal<I, I::Predicate>) -> QueryResult<I> {
         let Goal { param_env, predicate } = goal;
         let kind = predicate.kind();
         if let Some(kind) = kind.no_bound_vars() {

compiler/rustc_next_trait_solver/src/solve/search_graph.rs

@@ -1,13 +1,14 @@
 use std::convert::Infallible;
 use std::marker::PhantomData;
 
+use rustc_type_ir::data_structures::ensure_sufficient_stack;
 use rustc_type_ir::search_graph::{self, PathKind};
 use rustc_type_ir::solve::{CanonicalInput, Certainty, NoSolution, QueryResult};
 use rustc_type_ir::{Interner, TypingMode};
 
-use super::inspect::ProofTreeBuilder;
-use super::{FIXPOINT_STEP_LIMIT, has_no_inference_or_external_constraints};
 use crate::delegate::SolverDelegate;
+use crate::solve::inspect::ProofTreeBuilder;
+use crate::solve::{EvalCtxt, FIXPOINT_STEP_LIMIT, has_no_inference_or_external_constraints};
 
 /// 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.
@@ -80,8 +81,8 @@ where
 
     fn on_stack_overflow(
         cx: I,
-        inspect: &mut ProofTreeBuilder<D>,
         input: CanonicalInput<I>,
+        inspect: &mut ProofTreeBuilder<D>,
     ) -> QueryResult<I> {
         inspect.canonical_goal_evaluation_overflow();
         response_no_constraints(cx, input, Certainty::overflow(true))
@@ -106,6 +107,21 @@ where
         let certainty = from_result.unwrap().value.certainty;
         response_no_constraints(cx, for_input, certainty)
     }
+
+    fn compute_goal(
+        search_graph: &mut SearchGraph<D>,
+        cx: I,
+        input: CanonicalInput<I>,
+        inspect: &mut Self::ProofTreeBuilder,
+    ) -> QueryResult<I> {
+        ensure_sufficient_stack(|| {
+            EvalCtxt::enter_canonical(cx, search_graph, input, inspect, |ecx, goal| {
+                let result = ecx.compute_goal(goal);
+                ecx.inspect.query_result(result);
+                result
+            })
+        })
+    }
 }
 
 fn response_no_constraints<I: Interner>(

compiler/rustc_type_ir/src/search_graph/mod.rs

@@ -21,7 +21,7 @@ use std::marker::PhantomData;
 use derive_where::derive_where;
 #[cfg(feature = "nightly")]
 use rustc_macros::{Decodable_NoContext, Encodable_NoContext, HashStable_NoContext};
-use tracing::debug;
+use tracing::{debug, instrument};
 
 use crate::data_structures::HashMap;
 
@@ -56,7 +56,7 @@ pub trait Cx: Copy {
     fn evaluation_is_concurrent(&self) -> bool;
 }
 
-pub trait Delegate {
+pub trait Delegate: Sized {
     type Cx: Cx;
     /// Whether to use the provisional cache. Set to `false` by a fuzzer when
     /// validating the search graph.
@@ -94,8 +94,8 @@ pub trait Delegate {
     ) -> bool;
     fn on_stack_overflow(
         cx: Self::Cx,
-        inspect: &mut Self::ProofTreeBuilder,
         input: <Self::Cx as Cx>::Input,
+        inspect: &mut Self::ProofTreeBuilder,
     ) -> <Self::Cx as Cx>::Result;
     fn on_fixpoint_overflow(
         cx: Self::Cx,
@@ -108,6 +108,13 @@ pub trait Delegate {
         for_input: <Self::Cx as Cx>::Input,
         from_result: <Self::Cx as Cx>::Result,
     ) -> <Self::Cx as Cx>::Result;
+
+    fn compute_goal(
+        search_graph: &mut SearchGraph<Self>,
+        cx: Self::Cx,
+        input: <Self::Cx as Cx>::Input,
+        inspect: &mut Self::ProofTreeBuilder,
+    ) -> <Self::Cx as Cx>::Result;
 }
 
 /// In the initial iteration of a cycle, we do not yet have a provisional
@@ -589,15 +596,15 @@ impl<D: Delegate<Cx = X>, X: Cx> SearchGraph<D> {
 
     /// Probably the most involved method of the whole solver.
     ///
-    /// Given some goal which is proven via the `prove_goal` closure, this
-    /// handles caching, overflow, and coinductive cycles.
-    pub fn with_new_goal(
+    /// While goals get computed via `D::compute_goal`, this function handles
+    /// caching, overflow, and cycles.
+    #[instrument(level = "debug", skip(self, cx, inspect), ret)]
+    pub fn evaluate_goal(
         &mut self,
         cx: X,
         input: X::Input,
         step_kind_from_parent: PathKind,
         inspect: &mut D::ProofTreeBuilder,
-        evaluate_goal: impl Fn(&mut Self, X, X::Input, &mut D::ProofTreeBuilder) -> X::Result + Copy,
     ) -> X::Result {
         let Some(available_depth) =
             AvailableDepth::allowed_depth_for_nested::<D>(self.root_depth, &self.stack)
@@ -666,7 +673,7 @@ impl<D: Delegate<Cx = X>, X: Cx> SearchGraph<D> {
         // must not be added to the global cache. Notably, this is the case for
         // trait solver cycles participants.
         let (evaluation_result, dep_node) =
-            cx.with_cached_task(|| self.evaluate_goal_in_task(cx, input, inspect, evaluate_goal));
+            cx.with_cached_task(|| self.evaluate_goal_in_task(cx, input, inspect));
 
         // We've finished computing the goal and have popped it from the stack,
         // lazily update its parent goal.
@@ -736,7 +743,7 @@ impl<D: Delegate<Cx = X>, X: Cx> SearchGraph<D> {
         }
 
         debug!("encountered stack overflow");
-        D::on_stack_overflow(cx, inspect, input)
+        D::on_stack_overflow(cx, input, inspect)
     }
 
     /// When reevaluating a goal with a changed provisional result, all provisional cache entry
@@ -1064,7 +1071,6 @@ impl<D: Delegate<Cx = X>, X: Cx> SearchGraph<D> {
         cx: X,
         input: X::Input,
         inspect: &mut D::ProofTreeBuilder,
-        evaluate_goal: impl Fn(&mut Self, X, X::Input, &mut D::ProofTreeBuilder) -> X::Result + Copy,
     ) -> EvaluationResult<X> {
         // We reset `encountered_overflow` each time we rerun this goal
         // but need to make sure we currently propagate it to the global
@@ -1073,7 +1079,7 @@ impl<D: Delegate<Cx = X>, X: Cx> SearchGraph<D> {
         let mut encountered_overflow = false;
         let mut i = 0;
         loop {
-            let result = evaluate_goal(self, cx, input, inspect);
+            let result = D::compute_goal(self, cx, input, inspect);
             let stack_entry = self.stack.pop();
             encountered_overflow |= stack_entry.encountered_overflow;
             debug_assert_eq!(stack_entry.input, input);