Auto merge of #160479 - nnethercote:opt-try_evaluate_obligations, r=jdonszelmann

Optimize `try_evaluate_obligations`



This function is very sub-optimal, perf-wise: it takes `self.obligations.pending` (with `mem::take`) and iterates over the elements, checking each one. But most of the time no progress is made and all the obligations get pushed back onto `self.obligations.pending`. This drain + reconstruct approach is very expensive, mostly because the new `pending` vec is built by pushing one element at a time, which requires repeated reallocations. And this vec can have thousands of elements in it, in extreme cases.

Also, `obligation` and `stalled_on` get passed by value to `evaluate_root_goal` (`obligation` as `goal`), which then usually passes the values back in the `GoalEvaluation` which is immediately deconstructed. This is a lot of wasted value moves.

This commit optimizes things in two ways.

- It prioritizes the hot path. This involves checking in advance if there is an inspector (usually not) and adding `goal_remains_stalled` which takes `stalled_on` by reference. This hot path avoids all the value moves and `GoalEvaluation` construction/deconstruction and gets to the very common "nothing needed to be done" outcome as quickly as possible.

- It uses `retain_mut` to update `self.obligations.pending`. This requires some adjustments (e.g. handling recursion via the `overflowed` flag with some cleanup code after the `retain_mut` call, and cloning obligations in the error cases).

r? @lcnr
cc @jdonszelmann @WaffleLapkin

Author: bors

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

@@ -179,6 +179,11 @@ pub trait SolverDelegateEvalExt: SolverDelegate {
         stalled_on: Option<GoalStalledOn<Self::Interner>>,
     ) -> Result<GoalEvaluation<Self::Interner>, NoSolution>;
 
+    /// Checks whether a stalled goal would remain stalled if re-evaluated, without consuming
+    /// `stalled_on`.
+    fn goal_remains_stalled(&self, stalled_on: &GoalStalledOn<Self::Interner>)
+    -> Option<Certainty>;
+
     /// Checks whether evaluating `goal` may hold while treating not-yet-defined
     /// opaque types as being kind of rigid.
     ///
@@ -260,6 +265,16 @@ where
         }
     }
 
+    fn goal_remains_stalled(
+        &self,
+        stalled_on: &GoalStalledOn<Self::Interner>,
+    ) -> Option<Certainty> {
+        match rerunning_stalled_goal_may_make_progress(self, Some(stalled_on)) {
+            RerunStalled::WontMakeProgress(certainty) => Some(certainty),
+            RerunStalled::MayMakeProgress => None,
+        }
+    }
+
     #[instrument(level = "debug", skip(self), ret)]
     fn root_goal_may_hold_opaque_types_jank(
         &self,

compiler/rustc_trait_selection/src/solve/fulfill.rs

@@ -138,7 +138,6 @@ impl<'tcx, E: 'tcx> FulfillmentCtxt<'tcx, E> {
     }
 
     fn inspect_evaluated_obligation(
-        &self,
         infcx: &InferCtxt<'tcx>,
         obligation: &PredicateObligation<'tcx>,
         result: &Result<GoalEvaluation<TyCtxt<'tcx>>, NoSolution>,
@@ -196,22 +195,39 @@ where
     fn try_evaluate_obligations(&mut self, infcx: &InferCtxt<'tcx>) -> TraitErrors<E> {
         assert_eq!(self.usable_in_snapshot, infcx.num_open_snapshots());
         let mut errors = TraitErrors::NoErrors;
+        let delegate = <&SolverDelegate<'tcx>>::from(infcx);
         loop {
             let mut any_changed = false;
-            for (mut obligation, stalled_on) in mem::take(&mut self.obligations.pending) {
-                let goal = obligation.as_goal();
-                let delegate = <&SolverDelegate<'tcx>>::from(infcx);
+            let mut overflowed = false;
+
+            self.obligations.pending.retain_mut(|(obligation, opt_stalled_on)| {
+                if overflowed {
+                    return false;
+                }
 
-                let result = delegate.evaluate_root_goal(goal, obligation.cause.span, stalled_on);
-                self.inspect_evaluated_obligation(infcx, &obligation, &result);
+                // Common case: still stalled; keep the obligation. This path is extremely hot in
+                // some cases; there can be thousands of pending obligations.
+                if let Some(stalled_on) = opt_stalled_on
+                    && let Some(certainty) = delegate.goal_remains_stalled(stalled_on)
+                    && matches!(certainty, Certainty::Maybe(_))
+                {
+                    return true;
+                }
+
+                let result = delegate.evaluate_root_goal(
+                    obligation.as_goal(),
+                    obligation.cause.span,
+                    opt_stalled_on.take(),
+                );
+                Self::inspect_evaluated_obligation(infcx, &obligation, &result);
                 let GoalEvaluation { goal, certainty, has_changed, stalled_on } = match result {
                     Ok(result) => result,
                     Err(NoSolution) => {
                         errors.push(E::from_solver_error(
                             infcx,
-                            NextSolverError::TrueError(obligation),
+                            NextSolverError::TrueError(obligation.clone()),
                         ));
-                        continue;
+                        return false;
                     }
                 };
 
@@ -229,9 +245,11 @@ where
                     obligation.recursion_depth += 1;
 
                     if !infcx.tcx.recursion_limit().value_within_limit(obligation.recursion_depth) {
-                        self.obligations.on_fulfillment_overflow(infcx);
-                        // Only return true errors that we have accumulated while processing.
-                        return errors;
+                        // At this point we want to stop evaluating goals. We can't break out of
+                        // `retain_mut`, so instead we set this flag which causes all other
+                        // elements to be skipped.
+                        overflowed = true;
+                        return false;
                     } else {
                         any_changed = true;
                     }
@@ -253,11 +271,24 @@ where
                         if infcx.in_hir_typeck
                             && (obligation.has_non_region_infer() || obligation.has_free_regions())
                         {
-                            infcx.push_hir_typeck_potentially_region_dependent_goal(obligation);
+                            infcx.push_hir_typeck_potentially_region_dependent_goal(
+                                obligation.clone(),
+                            );
                         }
+                        false
+                    }
+                    Certainty::Maybe(_) => {
+                        // Update `opt_stalled_on` goal, for the next retain_mut, because we are
+                        // running until a fixpoint.
+                        *opt_stalled_on = stalled_on;
+                        true
                     }
-                    Certainty::Maybe(_) => self.obligations.register(obligation, stalled_on),
                 }
+            });
+            if overflowed {
+                self.obligations.on_fulfillment_overflow(infcx);
+                // Only return true errors that we have accumulated while processing.
+                return errors;
             }
 
             if !any_changed {