handle overflow in the `EvalCtxt` separately

6 files changed, 156 insertions, 189 deletions

diff --git a/compiler/rustc_trait_selection/src/solve/assembly/mod.rs b/compiler/rustc_trait_selection/src/solve/assembly/mod.rs
index e57188edc3d..267f345c062 100644
--- a/compiler/rustc_trait_selection/src/solve/assembly/mod.rs
+++ b/compiler/rustc_trait_selection/src/solve/assembly/mod.rs
@@ -1,6 +1,5 @@
 //! Code shared by trait and projection goals for candidate assembly.
 
-use super::search_graph::OverflowHandler;
 use super::{EvalCtxt, SolverMode};
 use crate::traits::coherence;
 use rustc_hir::def_id::DefId;
@@ -315,7 +314,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             return ambig;
         }
 
-        let mut candidates = self.assemble_candidates_via_self_ty(goal);
+        let mut candidates = self.assemble_candidates_via_self_ty(goal, 0);
 
         self.assemble_blanket_impl_candidates(goal, &mut candidates);
 
@@ -351,6 +350,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
     fn assemble_candidates_via_self_ty<G: GoalKind<'tcx>>(
         &mut self,
         goal: Goal<'tcx, G>,
+        num_steps: usize,
     ) -> Vec<Candidate<'tcx>> {
         debug_assert_eq!(goal, self.resolve_vars_if_possible(goal));
         if let Some(ambig) = self.assemble_self_ty_infer_ambiguity_response(goal) {
@@ -369,7 +369,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
 
         self.assemble_coherence_unknowable_candidates(goal, &mut candidates);
 
-        self.assemble_candidates_after_normalizing_self_ty(goal, &mut candidates);
+        self.assemble_candidates_after_normalizing_self_ty(goal, &mut candidates, num_steps);
 
         candidates
     }
@@ -393,46 +393,40 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         &mut self,
         goal: Goal<'tcx, G>,
         candidates: &mut Vec<Candidate<'tcx>>,
+        num_steps: usize,
     ) {
         let tcx = self.tcx();
         let &ty::Alias(_, projection_ty) = goal.predicate.self_ty().kind() else { return };
 
-        let normalized_self_candidates: Result<_, NoSolution> =
-            self.probe(|_| CandidateKind::NormalizedSelfTyAssembly).enter(|ecx| {
-                ecx.with_incremented_depth(
-                    |ecx| {
-                        let result = ecx.evaluate_added_goals_and_make_canonical_response(
-                            Certainty::OVERFLOW,
-                        )?;
-                        Ok(vec![Candidate {
-                            source: CandidateSource::BuiltinImpl(BuiltinImplSource::Misc),
-                            result,
-                        }])
-                    },
-                    |ecx| {
-                        let normalized_ty = ecx.next_ty_infer();
-                        let normalizes_to_goal = goal.with(
-                            tcx,
-                            ty::ProjectionPredicate { projection_ty, term: normalized_ty.into() },
-                        );
-                        ecx.add_goal(normalizes_to_goal);
-                        let _ = ecx.try_evaluate_added_goals().inspect_err(|_| {
-                            debug!("self type normalization failed");
-                        })?;
-                        let normalized_ty = ecx.resolve_vars_if_possible(normalized_ty);
-                        debug!(?normalized_ty, "self type normalized");
-                        // NOTE: Alternatively we could call `evaluate_goal` here and only
-                        // have a `Normalized` candidate. This doesn't work as long as we
-                        // use `CandidateSource` in winnowing.
-                        let goal = goal.with(tcx, goal.predicate.with_self_ty(tcx, normalized_ty));
-                        Ok(ecx.assemble_candidates_via_self_ty(goal))
-                    },
-                )
-            });
-
-        if let Ok(normalized_self_candidates) = normalized_self_candidates {
-            candidates.extend(normalized_self_candidates);
-        }
+        candidates.extend(self.probe(|_| CandidateKind::NormalizedSelfTyAssembly).enter(|ecx| {
+            if num_steps < ecx.local_overflow_limit() {
+                let normalized_ty = ecx.next_ty_infer();
+                let normalizes_to_goal = goal.with(
+                    tcx,
+                    ty::ProjectionPredicate { projection_ty, term: normalized_ty.into() },
+                );
+                ecx.add_goal(normalizes_to_goal);
+                if let Err(NoSolution) = ecx.try_evaluate_added_goals() {
+                    debug!("self type normalization failed");
+                    return vec![];
+                }
+                let normalized_ty = ecx.resolve_vars_if_possible(normalized_ty);
+                debug!(?normalized_ty, "self type normalized");
+                // NOTE: Alternatively we could call `evaluate_goal` here and only
+                // have a `Normalized` candidate. This doesn't work as long as we
+                // use `CandidateSource` in winnowing.
+                let goal = goal.with(tcx, goal.predicate.with_self_ty(tcx, normalized_ty));
+                ecx.assemble_candidates_via_self_ty(goal, num_steps + 1)
+            } else {
+                match ecx.evaluate_added_goals_and_make_canonical_response(Certainty::OVERFLOW) {
+                    Ok(result) => vec![Candidate {
+                        source: CandidateSource::BuiltinImpl(BuiltinImplSource::Misc),
+                        result,
+                    }],
+                    Err(NoSolution) => vec![],
+                }
+            }
+        }));
     }
 
     #[instrument(level = "debug", skip_all)]
@@ -530,7 +524,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             ty::Alias(_, _) | ty::Placeholder(..) | ty::Error(_) => (),
 
             // FIXME: These should ideally not exist as a self type. It would be nice for
-            // the builtin auto trait impls of generators should instead directly recurse
+            // the builtin auto trait impls of generators to instead directly recurse
             // into the witness.
             ty::GeneratorWitness(_) | ty::GeneratorWitnessMIR(_, _) => (),
 
diff --git a/compiler/rustc_trait_selection/src/solve/eval_ctxt.rs b/compiler/rustc_trait_selection/src/solve/eval_ctxt.rs
index 7283334aab2..8a20552d339 100644
--- a/compiler/rustc_trait_selection/src/solve/eval_ctxt.rs
+++ b/compiler/rustc_trait_selection/src/solve/eval_ctxt.rs
@@ -15,7 +15,7 @@ use rustc_middle::traits::solve::{
     CanonicalInput, CanonicalResponse, Certainty, IsNormalizesToHack, PredefinedOpaques,
     PredefinedOpaquesData, QueryResult,
 };
-use rustc_middle::traits::DefiningAnchor;
+use rustc_middle::traits::{specialization_graph, DefiningAnchor};
 use rustc_middle::ty::{
     self, OpaqueTypeKey, Ty, TyCtxt, TypeFoldable, TypeSuperVisitable, TypeVisitable,
     TypeVisitableExt, TypeVisitor,
@@ -25,11 +25,10 @@ use rustc_span::DUMMY_SP;
 use std::io::Write;
 use std::ops::ControlFlow;
 
-use crate::traits::specialization_graph;
 use crate::traits::vtable::{count_own_vtable_entries, prepare_vtable_segments, VtblSegment};
 
 use super::inspect::ProofTreeBuilder;
-use super::search_graph::{self, OverflowHandler};
+use super::search_graph;
 use super::SolverMode;
 use super::{search_graph::SearchGraph, Goal};
 pub use select::InferCtxtSelectExt;
@@ -175,6 +174,10 @@ impl<'a, 'tcx> EvalCtxt<'a, 'tcx> {
         self.search_graph.solver_mode()
     }
 
+    pub(super) fn local_overflow_limit(&self) -> usize {
+        self.search_graph.local_overflow_limit()
+    }
+
     /// Creates a root evaluation context and search graph. This should only be
     /// used from outside of any evaluation, and other methods should be preferred
     /// over using this manually (such as [`InferCtxtEvalExt::evaluate_root_goal`]).
@@ -479,101 +482,22 @@ impl<'a, 'tcx> EvalCtxt<'a, 'tcx> {
         let inspect = self.inspect.new_evaluate_added_goals();
         let inspect = core::mem::replace(&mut self.inspect, inspect);
 
-        let mut goals = core::mem::replace(&mut self.nested_goals, NestedGoals::new());
-        let mut new_goals = NestedGoals::new();
-
-        let response = self.repeat_while_none(
-            |_| Ok(Certainty::OVERFLOW),
-            |this| {
-                this.inspect.evaluate_added_goals_loop_start();
-
-                let mut has_changed = Err(Certainty::Yes);
-
-                if let Some(goal) = goals.normalizes_to_hack_goal.take() {
-                    // Replace the goal with an unconstrained infer var, so the
-                    // RHS does not affect projection candidate assembly.
-                    let unconstrained_rhs = this.next_term_infer_of_kind(goal.predicate.term);
-                    let unconstrained_goal = goal.with(
-                        this.tcx(),
-                        ty::ProjectionPredicate {
-                            projection_ty: goal.predicate.projection_ty,
-                            term: unconstrained_rhs,
-                        },
-                    );
-
-                    let (_, certainty, instantiate_goals) =
-                        match this.evaluate_goal(IsNormalizesToHack::Yes, unconstrained_goal) {
-                            Ok(r) => r,
-                            Err(NoSolution) => return Some(Err(NoSolution)),
-                        };
-                    new_goals.goals.extend(instantiate_goals);
-
-                    // Finally, equate the goal's RHS with the unconstrained var.
-                    // We put the nested goals from this into goals instead of
-                    // next_goals to avoid needing to process the loop one extra
-                    // time if this goal returns something -- I don't think this
-                    // matters in practice, though.
-                    match this.eq_and_get_goals(
-                        goal.param_env,
-                        goal.predicate.term,
-                        unconstrained_rhs,
-                    ) {
-                        Ok(eq_goals) => {
-                            goals.goals.extend(eq_goals);
-                        }
-                        Err(NoSolution) => return Some(Err(NoSolution)),
-                    };
-
-                    // We only look at the `projection_ty` part here rather than
-                    // looking at the "has changed" return from evaluate_goal,
-                    // because we expect the `unconstrained_rhs` part of the predicate
-                    // to have changed -- that means we actually normalized successfully!
-                    if goal.predicate.projection_ty
-                        != this.resolve_vars_if_possible(goal.predicate.projection_ty)
-                    {
-                        has_changed = Ok(())
-                    }
-
-                    match certainty {
-                        Certainty::Yes => {}
-                        Certainty::Maybe(_) => {
-                            // We need to resolve vars here so that we correctly
-                            // deal with `has_changed` in the next iteration.
-                            new_goals.normalizes_to_hack_goal =
-                                Some(this.resolve_vars_if_possible(goal));
-                            has_changed = has_changed.map_err(|c| c.unify_with(certainty));
-                        }
-                    }
-                }
-
-                for goal in goals.goals.drain(..) {
-                    let (changed, certainty, instantiate_goals) =
-                        match this.evaluate_goal(IsNormalizesToHack::No, goal) {
-                            Ok(result) => result,
-                            Err(NoSolution) => return Some(Err(NoSolution)),
-                        };
-                    new_goals.goals.extend(instantiate_goals);
-
-                    if changed {
-                        has_changed = Ok(());
-                    }
-
-                    match certainty {
-                        Certainty::Yes => {}
-                        Certainty::Maybe(_) => {
-                            new_goals.goals.push(goal);
-                            has_changed = has_changed.map_err(|c| c.unify_with(certainty));
-                        }
-                    }
+        let mut response = Ok(Certainty::OVERFLOW);
+        for _ in 0..self.local_overflow_limit() {
+            // FIXME: This match is a bit ugly, it might be nice to change the inspect
+            // stuff to use a closure instead. which should hopefully simplify this a bit.
+            match self.evaluate_added_goals_step() {
+                Ok(Some(cert)) => {
+                    response = Ok(cert);
+                    break;
                 }
-
-                core::mem::swap(&mut new_goals, &mut goals);
-                match has_changed {
-                    Ok(()) => None,
-                    Err(certainty) => Some(Ok(certainty)),
+                Ok(None) => {}
+                Err(NoSolution) => {
+                    response = Err(NoSolution);
+                    break;
                 }
-            },
-        );
+            }
+        }
 
         self.inspect.eval_added_goals_result(response);
 
@@ -584,9 +508,84 @@ impl<'a, 'tcx> EvalCtxt<'a, 'tcx> {
         let goal_evaluations = std::mem::replace(&mut self.inspect, inspect);
         self.inspect.added_goals_evaluation(goal_evaluations);
 
-        self.nested_goals = goals;
         response
     }
+
+    /// Iterate over all added goals: returning `Ok(Some(_))` in case we can stop rerunning.
+    ///
+    /// Goals for the next step get directly added the the nested goals of the `EvalCtxt`.
+    fn evaluate_added_goals_step(&mut self) -> Result<Option<Certainty>, NoSolution> {
+        let tcx = self.tcx();
+        let mut goals = core::mem::replace(&mut self.nested_goals, NestedGoals::new());
+
+        self.inspect.evaluate_added_goals_loop_start();
+        // If this loop did not result in any progress, what's our final certainty.
+        let mut unchanged_certainty = Some(Certainty::Yes);
+        if let Some(goal) = goals.normalizes_to_hack_goal.take() {
+            // Replace the goal with an unconstrained infer var, so the
+            // RHS does not affect projection candidate assembly.
+            let unconstrained_rhs = self.next_term_infer_of_kind(goal.predicate.term);
+            let unconstrained_goal = goal.with(
+                tcx,
+                ty::ProjectionPredicate {
+                    projection_ty: goal.predicate.projection_ty,
+                    term: unconstrained_rhs,
+                },
+            );
+
+            let (_, certainty, instantiate_goals) =
+                self.evaluate_goal(IsNormalizesToHack::Yes, unconstrained_goal)?;
+            self.add_goals(instantiate_goals);
+
+            // Finally, equate the goal's RHS with the unconstrained var.
+            // We put the nested goals from this into goals instead of
+            // next_goals to avoid needing to process the loop one extra
+            // time if this goal returns something -- I don't think this
+            // matters in practice, though.
+            let eq_goals =
+                self.eq_and_get_goals(goal.param_env, goal.predicate.term, unconstrained_rhs)?;
+            goals.goals.extend(eq_goals);
+
+            // We only look at the `projection_ty` part here rather than
+            // looking at the "has changed" return from evaluate_goal,
+            // because we expect the `unconstrained_rhs` part of the predicate
+            // to have changed -- that means we actually normalized successfully!
+            if goal.predicate.projection_ty
+                != self.resolve_vars_if_possible(goal.predicate.projection_ty)
+            {
+                unchanged_certainty = None;
+            }
+
+            match certainty {
+                Certainty::Yes => {}
+                Certainty::Maybe(_) => {
+                    // We need to resolve vars here so that we correctly
+                    // deal with `has_changed` in the next iteration.
+                    self.set_normalizes_to_hack_goal(self.resolve_vars_if_possible(goal));
+                    unchanged_certainty = unchanged_certainty.map(|c| c.unify_with(certainty));
+                }
+            }
+        }
+
+        for goal in goals.goals.drain(..) {
+            let (has_changed, certainty, instantiate_goals) =
+                self.evaluate_goal(IsNormalizesToHack::No, goal)?;
+            self.add_goals(instantiate_goals);
+            if has_changed {
+                unchanged_certainty = None;
+            }
+
+            match certainty {
+                Certainty::Yes => {}
+                Certainty::Maybe(_) => {
+                    self.add_goal(goal);
+                    unchanged_certainty = unchanged_certainty.map(|c| c.unify_with(certainty));
+                }
+            }
+        }
+
+        Ok(unchanged_certainty)
+    }
 }
 
 impl<'tcx> EvalCtxt<'_, 'tcx> {
diff --git a/compiler/rustc_trait_selection/src/solve/eval_ctxt/select.rs b/compiler/rustc_trait_selection/src/solve/eval_ctxt/select.rs
index 2b46e5b8b01..ca4a4c9510c 100644
--- a/compiler/rustc_trait_selection/src/solve/eval_ctxt/select.rs
+++ b/compiler/rustc_trait_selection/src/solve/eval_ctxt/select.rs
@@ -14,7 +14,6 @@ use rustc_span::DUMMY_SP;
 use crate::solve::assembly::{Candidate, CandidateSource};
 use crate::solve::eval_ctxt::{EvalCtxt, GenerateProofTree};
 use crate::solve::inspect::ProofTreeBuilder;
-use crate::solve::search_graph::OverflowHandler;
 use crate::traits::StructurallyNormalizeExt;
 use crate::traits::TraitEngineExt;
 
@@ -143,7 +142,7 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
         // the cycle anyways one step later.
         EvalCtxt::enter_canonical(
             self.tcx(),
-            self.search_graph(),
+            self.search_graph,
             canonical_input,
             // FIXME: This is wrong, idk if we even want to track stuff here.
             &mut ProofTreeBuilder::new_noop(),
diff --git a/compiler/rustc_trait_selection/src/solve/search_graph/mod.rs b/compiler/rustc_trait_selection/src/solve/search_graph/mod.rs
index b70ec48ee2f..f84445e8820 100644
--- a/compiler/rustc_trait_selection/src/solve/search_graph/mod.rs
+++ b/compiler/rustc_trait_selection/src/solve/search_graph/mod.rs
@@ -27,6 +27,7 @@ struct StackElem<'tcx> {
 
 pub(super) struct SearchGraph<'tcx> {
     mode: SolverMode,
+    local_overflow_limit: usize,
     /// The stack of goals currently being computed.
     ///
     /// An element is *deeper* in the stack if its index is *lower*.
@@ -39,6 +40,7 @@ impl<'tcx> SearchGraph<'tcx> {
     pub(super) fn new(tcx: TyCtxt<'tcx>, mode: SolverMode) -> SearchGraph<'tcx> {
         Self {
             mode,
+            local_overflow_limit: tcx.recursion_limit().0.ilog2() as usize,
             stack: Default::default(),
             overflow_data: OverflowData::new(tcx),
             provisional_cache: ProvisionalCache::empty(),
@@ -49,6 +51,10 @@ impl<'tcx> SearchGraph<'tcx> {
         self.mode
     }
 
+    pub(super) fn local_overflow_limit(&self) -> usize {
+        self.local_overflow_limit
+    }
+
     /// We do not use the global cache during coherence.
     ///
     /// The trait solver behavior is different for coherence
diff --git a/compiler/rustc_trait_selection/src/solve/search_graph/overflow.rs b/compiler/rustc_trait_selection/src/solve/search_graph/overflow.rs
index f54563557d7..4abde2c96d4 100644
--- a/compiler/rustc_trait_selection/src/solve/search_graph/overflow.rs
+++ b/compiler/rustc_trait_selection/src/solve/search_graph/overflow.rs
@@ -5,7 +5,7 @@ use rustc_middle::ty::TyCtxt;
 use rustc_session::Limit;
 
 use super::SearchGraph;
-use crate::solve::{response_no_constraints_raw, EvalCtxt};
+use crate::solve::response_no_constraints_raw;
 
 /// When detecting a solver overflow, we return ambiguity. Overflow can be
 /// *hidden* by either a fatal error in an **AND** or a trivial success in an **OR**.
@@ -73,33 +73,6 @@ pub(in crate::solve) trait OverflowHandler<'tcx> {
         self.search_graph().overflow_data.deal_with_overflow();
         on_overflow(self)
     }
-
-    // Increment the `additional_depth` by one and evaluate `body`, or `on_overflow`
-    // if the depth is overflown.
-    fn with_incremented_depth<T>(
-        &mut self,
-        on_overflow: impl FnOnce(&mut Self) -> T,
-        body: impl FnOnce(&mut Self) -> T,
-    ) -> T {
-        let depth = self.search_graph().stack.len();
-        self.search_graph().overflow_data.additional_depth += 1;
-
-        let result = if self.search_graph().overflow_data.has_overflow(depth) {
-            self.search_graph().overflow_data.deal_with_overflow();
-            on_overflow(self)
-        } else {
-            body(self)
-        };
-
-        self.search_graph().overflow_data.additional_depth -= 1;
-        result
-    }
-}
-
-impl<'tcx> OverflowHandler<'tcx> for EvalCtxt<'_, 'tcx> {
-    fn search_graph(&mut self) -> &mut SearchGraph<'tcx> {
-        &mut self.search_graph
-    }
 }
 
 impl<'tcx> OverflowHandler<'tcx> for SearchGraph<'tcx> {
diff --git a/compiler/rustc_trait_selection/src/solve/trait_goals.rs b/compiler/rustc_trait_selection/src/solve/trait_goals.rs
index 6e552633246..6a96ea38555 100644
--- a/compiler/rustc_trait_selection/src/solve/trait_goals.rs
+++ b/compiler/rustc_trait_selection/src/solve/trait_goals.rs
@@ -1,7 +1,6 @@
 //! Dealing with trait goals, i.e. `T: Trait<'a, U>`.
 
 use super::assembly::{self, structural_traits};
-use super::search_graph::OverflowHandler;
 use super::{EvalCtxt, SolverMode};
 use rustc_hir::def_id::DefId;
 use rustc_hir::{LangItem, Movability};
@@ -874,7 +873,9 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
     }
 
     /// Normalize a non-self type when it is structually matched on when solving
-    /// a built-in goal. This is handled already through `assemble_candidates_after_normalizing_self_ty`
+    /// a built-in goal.
+    ///
+    /// This is handled already through `assemble_candidates_after_normalizing_self_ty`
     /// for the self type, but for other goals, additional normalization of other
     /// arguments may be needed to completely implement the semantics of the trait.
     ///
@@ -889,27 +890,22 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             return Ok(Some(ty));
         }
 
-        self.repeat_while_none(
-            |_| Ok(None),
-            |ecx| {
-                let ty::Alias(_, projection_ty) = *ty.kind() else {
-                    return Some(Ok(Some(ty)));
-                };
+        for _ in 0..self.local_overflow_limit() {
+            let ty::Alias(_, projection_ty) = *ty.kind() else {
+                return Ok(Some(ty));
+            };
 
-                let normalized_ty = ecx.next_ty_infer();
-                let normalizes_to_goal = Goal::new(
-                    ecx.tcx(),
-                    param_env,
-                    ty::ProjectionPredicate { projection_ty, term: normalized_ty.into() },
-                );
-                ecx.add_goal(normalizes_to_goal);
-                if let Err(err) = ecx.try_evaluate_added_goals() {
-                    return Some(Err(err));
-                }
+            let normalized_ty = self.next_ty_infer();
+            let normalizes_to_goal = Goal::new(
+                self.tcx(),
+                param_env,
+                ty::ProjectionPredicate { projection_ty, term: normalized_ty.into() },
+            );
+            self.add_goal(normalizes_to_goal);
+            self.try_evaluate_added_goals()?;
+            ty = self.resolve_vars_if_possible(normalized_ty);
+        }
 
-                ty = ecx.resolve_vars_if_possible(normalized_ty);
-                None
-            },
-        )
+        Ok(None)
     }
 }