fix: preserve frontend fulfillment overflow behavior

Signed-off-by: Amirhossein Akhlaghpour <m9.akhlaghpoor@gmail.com>

Author: amirHdev

compiler/rustc_next_trait_solver/src/solve/fulfill.rs

@@ -15,11 +15,30 @@ pub enum NextSolverError<O> {
     Overflow(O),
 }
 
+/// Controls when fulfillment detects recursion overflow.
+///
+/// rustc currently checks after a goal makes inference progress, while
+/// rust-analyzer checks before evaluating an obligation which has already
+/// reached the recursion limit.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub enum FulfillmentOverflowMode {
+    /// Check after a goal makes inference progress.
+    AfterProgress,
+
+    /// Check before evaluating an obligation at the recursion limit.
+    BeforeEvaluation,
+}
+
+// FIXME: Do we need to use a `ThinVec` here?
 type PendingObligations<I, O> = ThinVec<(O, Option<GoalStalledOn<I>>)>;
 
 #[derive(Debug)]
 struct ObligationStorage<I: Interner, O> {
     /// Obligations which resulted in overflow in fulfillment itself.
+    ///
+    /// We cannot eagerly return these as errors, so we instead store them here
+    /// to avoid recomputing them each time `try_evaluate_obligations` is called.
+    /// This also allows the frontend to construct the correct error for them.
     overflowed: Vec<O>,
 
     pending: PendingObligations<I, O>,
@@ -96,6 +115,14 @@ impl<I: Interner, O> ObligationStorage<I, O> {
     }
 }
 
+/// A fulfillment engine using the new trait solver.
+///
+/// This is mostly identical to how `evaluate_all` works inside of the solver,
+/// except that it is possible to add new obligations later and the frontend
+/// needs to retain its obligation representation for diagnostics.
+///
+/// It is also likely that we want to use different data structures here, as
+/// fulfillment deals with far more root goals than `evaluate_all`.
 #[derive(Debug)]
 pub struct FulfillmentCtxt<I: Interner, O: FulfillmentObligation<I>> {
     obligations: ObligationStorage<I, O>,
@@ -119,6 +146,9 @@ impl<I: Interner, O: FulfillmentObligation<I>> FulfillmentCtxt<I, O> {
         if let Some(GoalEvaluation { certainty, stalled_on, .. }) =
             compute_goal_fast_path(delegate, obligation.as_goal(), obligation.span())
         {
+            // If we can take the fast path, do not add a successful goal to
+            // the pending obligations. For `Certainty::Maybe`, retain the
+            // precise `stalled_on` information for later re-evaluation.
             match certainty {
                 Certainty::Yes => {}
                 Certainty::Maybe(_) => {
@@ -135,6 +165,11 @@ impl<I: Interner, O: FulfillmentObligation<I>> FulfillmentCtxt<I, O> {
         D: SolverDelegate<Interner = I>,
     {
         delegate.probe(|| {
+            // IMPORTANT: we must not resolve any inference variables in the
+            // obligations, as this is all happening inside of a probe. The
+            // probe makes sure we collect every obligation involved in the
+            // overflow. Conceptually, we check which goals would change if we
+            // performed one more fulfillment iteration.
             let overflowed = self
                 .obligations
                 .pending
@@ -157,6 +192,7 @@ impl<I: Interner, O: FulfillmentObligation<I>> FulfillmentCtxt<I, O> {
     pub fn try_evaluate_obligations<D, Inspect, OnSuccess>(
         &mut self,
         delegate: &D,
+        overflow_mode: FulfillmentOverflowMode,
         mut inspect: Inspect,
         mut on_success: OnSuccess,
     ) -> Vec<NextSolverError<O>>
@@ -171,6 +207,13 @@ impl<I: Interner, O: FulfillmentObligation<I>> FulfillmentCtxt<I, O> {
             let mut any_changed = false;
 
             for (mut obligation, stalled_on) in std::mem::take(&mut self.obligations.pending) {
+                if overflow_mode == FulfillmentOverflowMode::BeforeEvaluation
+                    && obligation.recursion_depth() >= delegate.cx().recursion_limit()
+                {
+                    self.on_fulfillment_overflow(delegate);
+                    return errors;
+                }
+
                 let goal = obligation.as_goal();
                 let result = delegate.evaluate_root_goal(goal, obligation.span(), stalled_on);
 
@@ -184,15 +227,26 @@ impl<I: Interner, O: FulfillmentObligation<I>> FulfillmentCtxt<I, O> {
                     }
                 };
 
-                // Reuse the eagerly resolved predicate during the next iteration.
+                // We resolved the goal in `evaluate_root_goal`; retain the
+                // eagerly resolved predicate to avoid repeating this work in
+                // the next iteration. This does not resolve the inference
+                // variables constrained by evaluating the goal.
                 obligation.set_predicate(goal.predicate);
 
                 if has_changed == HasChanged::Yes {
+                    // Track the number of times this root goal has resulted in
+                    // inference progress. This does not precisely model the old
+                    // solver's recursion depth, as fulfillment only processes
+                    // root obligations, but it is a good approximation and
+                    // should only overflow in pathological cases.
                     let depth = obligation.recursion_depth() + 1;
                     obligation.set_recursion_depth(depth);
 
-                    if depth > delegate.cx().recursion_limit() {
+                    if overflow_mode == FulfillmentOverflowMode::AfterProgress
+                        && depth > delegate.cx().recursion_limit()
+                    {
                         self.on_fulfillment_overflow(delegate);
+                        // Only return true errors accumulated while processing.
                         return errors;
                     }
 

compiler/rustc_trait_selection/src/solve/rustc_fulfill.rs

@@ -7,7 +7,8 @@ use rustc_infer::traits::{
 };
 use rustc_middle::ty::{self, TyCtxt, TypeVisitableExt, TypingMode};
 use rustc_next_trait_solver::solve::fulfill::{
-    FulfillmentCtxt as SolverFulfillmentCtxt, NextSolverError as SolverNextSolverError,
+    FulfillmentCtxt as SolverFulfillmentCtxt, FulfillmentOverflowMode,
+    NextSolverError as SolverNextSolverError,
 };
 use rustc_next_trait_solver::solve::{GoalEvaluation, MaybeInfo, StalledOnCoroutines};
 use tracing::instrument;
@@ -28,7 +29,8 @@ pub struct FulfillmentCtxt<'tcx, E: 'tcx> {
 
     /// The snapshot in which this context was created. Using the context
     /// outside of this snapshot leads to subtle bugs if the snapshot
-    /// gets rolled back.
+    /// gets rolled back. Because of this we explicitly check that we only
+    /// use the context in exactly this snapshot.
     usable_in_snapshot: usize,
     _errors: PhantomData<E>,
 }
@@ -100,13 +102,22 @@ where
         self.core
             .try_evaluate_obligations(
                 delegate,
+                FulfillmentOverflowMode::AfterProgress,
                 |obligation, _, result| {
                     Self::inspect_evaluated_obligation(infcx, obligation, result);
                 },
                 |obligation| {
-                    // Goals may depend on structural identity. Region
-                    // uniquification at the start of MIR borrowck may cause
-                    // things to no longer be structurally identical.
+                    // Goals may depend on structural identity. Region uniquification at the
+                    // start of MIR borrowck may cause things to no longer be so, potentially
+                    // causing an ICE.
+                    //
+                    // While we uniquify root goals in HIR this does not handle cases where
+                    // regions are hidden inside of a type or const inference variable.
+                    //
+                    // FIXME(-Znext-solver): This does not handle inference variables hidden
+                    // inside of an opaque type, e.g. if there's `Opaque = (?x, ?x)` in the
+                    // storage, we can also rely on structural identity of `?x` even if we
+                    // later uniquify it in MIR borrowck.
                     if infcx.in_hir_typeck
                         && (obligation.has_non_region_infer() || obligation.has_free_regions())
                     {
@@ -133,6 +144,7 @@ where
         infcx: &InferCtxt<'tcx>,
         vid: ty::TyVid,
     ) -> PredicateObligations<'tcx> {
+        // `-Zdisable-fast-paths`: same gate as the other new-solver fast paths.
         if infcx.tcx.disable_trait_solver_fast_paths() {
             return self.pending_obligations();
         }
@@ -142,6 +154,14 @@ where
                 return true;
             };
 
+            // Don't reuse the sub-unification roots cached on `stalled_on`:
+            // a later sub-unification merge can have changed which root
+            // each stalled var belongs to, so the cached info can be stale.
+            // Walk `stalled_vars` and recompute the current root instead.
+            //
+            // Conservative here: if a stalled var no longer resolves to an
+            // infer var, some unification happened, so the goal is no longer
+            // stalled. Include it to be re-evaluated downstream.
             stalled_on.stalled_vars.iter().filter_map(|arg| arg.as_type()).any(|ty| {
                 match *infcx.shallow_resolve(ty).kind() {
                     ty::Infer(ty::TyVar(tv)) => infcx.sub_unification_table_root_var(tv) == vid,
@@ -155,6 +175,7 @@ where
         &self,
         infcx: &InferCtxt<'tcx>,
     ) -> PredicateObligations<'tcx> {
+        // `-Zdisable-fast-paths`: same gate as the other new-solver fast paths.
         if infcx.tcx.disable_trait_solver_fast_paths() {
             return self.pending_obligations();
         }
@@ -164,6 +185,9 @@ where
                 return true;
             };
 
+            // If the stalled vars don't have float infers, the nested goals
+            // won't have them either. We only create float infers for
+            // user-written literals.
             stalled_on
                 .stalled_vars
                 .iter()

compiler/rustc_type_ir/src/solve/fulfill.rs

@@ -3,8 +3,7 @@ use crate::Interner;
 
 /// An obligation that can be processed by the shared fulfillment engine.
 ///
-/// The concrete obligation and its diagnostic cause remain owned by the
-/// frontend. Fulfillment only needs the goal, span, and recursion depth.
+/// The shared engine only accesses the parts needed for fulfillment through this trait.
 pub trait FulfillmentObligation<I: Interner>: Clone {
     fn as_goal(&self) -> Goal<I, I::Predicate>;