Avoid redundant next-solver fulfillment scans

compiler/rustc_infer/src/infer/mod.rs

@@ -164,12 +164,12 @@ pub struct InferCtxtInner<'tcx> {
 }
 
 impl<'tcx> InferCtxtInner<'tcx> {
-    fn new() -> InferCtxtInner<'tcx> {
+    fn new(next_trait_solver: bool) -> InferCtxtInner<'tcx> {
         InferCtxtInner {
             undo_log: InferCtxtUndoLogs::default(),
 
             projection_cache: Default::default(),
-            type_variable_storage: Default::default(),
+            type_variable_storage: type_variable::TypeVariableStorage::new(next_trait_solver),
             const_unification_storage: Default::default(),
             int_unification_storage: Default::default(),
             float_unification_storage: Default::default(),
@@ -677,7 +677,7 @@ impl<'tcx> InferCtxtBuilder<'tcx> {
             considering_regions,
             in_hir_typeck,
             skip_leak_check,
-            inner: RefCell::new(InferCtxtInner::new()),
+            inner: RefCell::new(InferCtxtInner::new(next_trait_solver)),
             lexical_region_resolutions: RefCell::new(None),
             selection_cache: Default::default(),
             evaluation_cache: Default::default(),
@@ -1593,6 +1593,19 @@ impl<'tcx> InferCtxt<'tcx> {
         self.typing_env(param_env).as_query_input(value)
     }
 
+    #[inline]
+    pub fn stalled_goal_revisions(&self) -> (u64, u64) {
+        self.inner.borrow().type_variable_storage.stalled_goal_revisions()
+    }
+
+    #[inline]
+    pub fn stalled_goal_sub_var_is_root(&self, vid: ty::TyVid) -> bool {
+        let inner = self.inner.borrow();
+        let table = inner.type_variable_storage.sub_unification_table_ref();
+
+        (vid.as_u32() as usize) < table.len() && table.try_probe_value(vid).is_some()
+    }
+
     /// The returned function is used in a fast path. If it returns `true` the variable is
     /// unchanged, `false` indicates that the status is unknown.
     #[inline]

compiler/rustc_infer/src/infer/type_variable.rs

@@ -2,7 +2,7 @@ use std::cmp;
 use std::marker::PhantomData;
 use std::ops::Range;
 
-use rustc_data_structures::undo_log::Rollback;
+use rustc_data_structures::undo_log::{Rollback, UndoLogs};
 use rustc_data_structures::{snapshot_vec as sv, unify as ut};
 use rustc_hir::HirId;
 use rustc_hir::def_id::DefId;
@@ -19,6 +19,8 @@ use crate::infer::InferCtxtUndoLogs;
 pub(crate) enum UndoLog<'tcx> {
     EqRelation(sv::UndoLog<ut::Delegate<TyVidEqKey<'tcx>>>),
     SubRelation(sv::UndoLog<ut::Delegate<TyVidSubKey>>),
+    StalledGoalRevision(u64),
+    StalledGoalSubRevision(u64),
 }
 
 /// Convert from a specific kind of undo to the more general UndoLog
@@ -52,6 +54,12 @@ impl<'tcx> Rollback<UndoLog<'tcx>> for TypeVariableStorage<'tcx> {
         match undo {
             UndoLog::EqRelation(undo) => self.eq_relations.reverse(undo),
             UndoLog::SubRelation(undo) => self.sub_unification_table.reverse(undo),
+            UndoLog::StalledGoalRevision(revision) => {
+                self.stalled_goal_revision = revision;
+            }
+            UndoLog::StalledGoalSubRevision(revision) => {
+                self.stalled_goal_sub_revision = revision;
+            }
         }
     }
 }
@@ -83,6 +91,16 @@ pub(crate) struct TypeVariableStorage<'tcx> {
     /// type of `x` is only a supertype of the argument of `returns_arg`. We
     /// still want to suggest specifying the type of the argument.
     sub_unification_table: ut::UnificationTableStorage<TyVidSubKey>,
+
+    // Revisions for inference changes relevant to stalled type-variable
+    // goals. Revision changes participate in the inference undo log so
+    // rolling back a snapshot restores them together with inference state.
+    stalled_goal_revision: u64,
+    stalled_goal_sub_revision: u64,
+
+    /// Only track stalled-goal revisions for inference contexts using
+    /// the next trait solver.
+    track_stalled_goal_revisions: bool,
 }
 
 pub(crate) struct TypeVariableTable<'a, 'tcx> {
@@ -140,6 +158,10 @@ impl<'tcx> TypeVariableValue<'tcx> {
 }
 
 impl<'tcx> TypeVariableStorage<'tcx> {
+    pub(crate) fn new(track_stalled_goal_revisions: bool) -> Self {
+        Self { track_stalled_goal_revisions, ..Default::default() }
+    }
+
     #[inline]
     pub(crate) fn with_log<'a>(
         &'a mut self,
@@ -161,9 +183,28 @@ impl<'tcx> TypeVariableStorage<'tcx> {
     pub(crate) fn sub_unification_table_ref(&self) -> &ut::UnificationTableStorage<TyVidSubKey> {
         &self.sub_unification_table
     }
+
+    #[inline]
+    pub(crate) fn stalled_goal_revisions(&self) -> (u64, u64) {
+        (self.stalled_goal_revision, self.stalled_goal_sub_revision)
+    }
 }
 
 impl<'tcx> TypeVariableTable<'_, 'tcx> {
+    #[inline]
+    fn bump_stalled_goal_revision(&mut self) {
+        let revision = self.storage.stalled_goal_revision;
+        self.undo_log.push(UndoLog::StalledGoalRevision(revision));
+        self.storage.stalled_goal_revision = revision.wrapping_add(1);
+    }
+
+    #[inline]
+    fn bump_stalled_goal_sub_revision(&mut self) {
+        let revision = self.storage.stalled_goal_sub_revision;
+        self.undo_log.push(UndoLog::StalledGoalSubRevision(revision));
+        self.storage.stalled_goal_sub_revision = revision.wrapping_add(1);
+    }
+
     /// Returns the origin that was given when `vid` was created.
     ///
     /// Note that this function does not return care whether
@@ -178,6 +219,12 @@ impl<'tcx> TypeVariableTable<'_, 'tcx> {
     pub(crate) fn equate(&mut self, a: ty::TyVid, b: ty::TyVid) {
         debug_assert!(self.probe(a).is_unknown());
         debug_assert!(self.probe(b).is_unknown());
+
+        if self.storage.track_stalled_goal_revisions {
+            self.bump_stalled_goal_revision();
+            self.bump_stalled_goal_sub_revision();
+        }
+
         self.eq_relations().union(a, b);
         self.sub_unification_table().union(a, b);
     }
@@ -189,6 +236,11 @@ impl<'tcx> TypeVariableTable<'_, 'tcx> {
     pub(crate) fn sub_unify(&mut self, a: ty::TyVid, b: ty::TyVid) {
         debug_assert!(self.probe(a).is_unknown());
         debug_assert!(self.probe(b).is_unknown());
+
+        if self.storage.track_stalled_goal_revisions {
+            self.bump_stalled_goal_sub_revision();
+        }
+
         self.sub_unification_table().union(a, b);
     }
 
@@ -204,6 +256,11 @@ impl<'tcx> TypeVariableTable<'_, 'tcx> {
             "instantiating type variable `{vid:?}` twice: new-value = {ty:?}, old-value={:?}",
             self.eq_relations().probe_value(vid)
         );
+
+        if self.storage.track_stalled_goal_revisions {
+            self.bump_stalled_goal_revision();
+        }
+
         self.eq_relations().union_value(vid, TypeVariableValue::Known { value: ty });
     }
 

compiler/rustc_trait_selection/src/solve/fulfill.rs

@@ -1,6 +1,7 @@
 use std::marker::PhantomData;
 use std::mem;
 
+use rustc_data_structures::fx::FxIndexSet;
 use rustc_infer::infer::InferCtxt;
 use rustc_infer::traits::query::NoSolution;
 use rustc_infer::traits::{
@@ -9,7 +10,8 @@ use rustc_infer::traits::{
 use rustc_middle::ty::{self, TyCtxt, TypeVisitableExt, TypingMode};
 use rustc_next_trait_solver::solve::fast_path::compute_goal_fast_path;
 use rustc_next_trait_solver::solve::{
-    GoalEvaluation, GoalStalledOn, HasChanged, SolverDelegateEvalExt as _, StalledOnCoroutines,
+    GoalEvaluation, GoalStalledOn, GoalStalledOnOpaques, HasChanged, SolverDelegateEvalExt as _,
+    StalledOnCoroutines, TyOrConstInferVar,
 };
 use thin_vec::ThinVec;
 use tracing::instrument;
@@ -44,6 +46,22 @@ pub struct FulfillmentCtxt<'tcx, E: 'tcx> {
     /// gets rolled back. Because of this we explicitly check that we only
     /// use the context in exactly this snapshot.
     usable_in_snapshot: usize,
+
+    last_stalled_goal_revision: u64,
+    last_stalled_goal_sub_revision: u64,
+
+    /// Over-approximation of sub-unification roots referenced by
+    /// trackable stalled obligations.
+    stalled_sub_roots: FxIndexSet<ty::TyVid>,
+
+    /// Whether any trackable stalled obligation requires the opaque
+    /// type storage to remain empty.
+    stalled_on_empty_opaques: bool,
+
+    /// Whether every pending obligation can use the context-wide
+    /// stalled-goal fast path.
+    all_pending_trackable: bool,
+
     _errors: PhantomData<E>,
 }
 
@@ -129,13 +147,41 @@ impl<'tcx, E: 'tcx> FulfillmentCtxt<'tcx, E> {
             "new trait solver fulfillment context created when \
             infcx is set up for old trait solver"
         );
+        let (revision, sub_revision) = infcx.stalled_goal_revisions();
+
         FulfillmentCtxt {
             obligations: Default::default(),
             usable_in_snapshot: infcx.num_open_snapshots(),
+            last_stalled_goal_revision: revision,
+            last_stalled_goal_sub_revision: sub_revision,
+            stalled_sub_roots: Default::default(),
+            stalled_on_empty_opaques: false,
+            all_pending_trackable: true,
             _errors: PhantomData,
         }
     }
 
+    fn record_trackable_stalled_on(
+        stalled_on: &GoalStalledOn<TyCtxt<'tcx>>,
+        stalled_sub_roots: &mut FxIndexSet<ty::TyVid>,
+        stalled_on_empty_opaques: &mut bool,
+    ) -> bool {
+        if stalled_on.stalled_vars.iter().any(|var| !matches!(*var, TyOrConstInferVar::Ty(_))) {
+            return false;
+        }
+
+        match stalled_on.opaques {
+            GoalStalledOnOpaques::No => {}
+            GoalStalledOnOpaques::Yes { num_opaques_in_storage: 0, .. } => {
+                *stalled_on_empty_opaques = true;
+            }
+            GoalStalledOnOpaques::Yes { .. } => return false,
+        }
+
+        stalled_sub_roots.extend(stalled_on.sub_roots.iter().copied());
+        true
+    }
+
     fn inspect_evaluated_obligation(
         infcx: &InferCtxt<'tcx>,
         obligation: &PredicateObligation<'tcx>,
@@ -172,10 +218,23 @@ where
             match certainty {
                 Certainty::Yes => {}
                 Certainty::Maybe(_) => {
+                    if let Some(stalled_on) = &stalled_on {
+                        if !Self::record_trackable_stalled_on(
+                            stalled_on,
+                            &mut self.stalled_sub_roots,
+                            &mut self.stalled_on_empty_opaques,
+                        ) {
+                            self.all_pending_trackable = false;
+                        }
+                    } else {
+                        self.all_pending_trackable = false;
+                    }
+
                     self.obligations.register(obligation, stalled_on);
                 }
             }
         } else {
+            self.all_pending_trackable = false;
             self.obligations.register(obligation, None);
         }
     }
@@ -195,9 +254,45 @@ where
         assert_eq!(self.usable_in_snapshot, infcx.num_open_snapshots());
         let mut errors = TraitErrors::NoErrors;
         let delegate = <&SolverDelegate<'tcx>>::from(infcx);
+
+        let (revision, sub_revision) = infcx.stalled_goal_revisions();
+
+        if self.obligations.pending.is_empty() {
+            self.last_stalled_goal_revision = revision;
+            self.last_stalled_goal_sub_revision = sub_revision;
+            self.stalled_sub_roots.clear();
+            self.stalled_on_empty_opaques = false;
+            self.all_pending_trackable = true;
+            return errors;
+        }
+
+        if !infcx.tcx.disable_trait_solver_fast_paths()
+            && self.all_pending_trackable
+            && self.last_stalled_goal_revision == revision
+        {
+            let sub_roots_unchanged = self.last_stalled_goal_sub_revision == sub_revision
+                || self
+                    .stalled_sub_roots
+                    .iter()
+                    .all(|&vid| infcx.stalled_goal_sub_var_is_root(vid));
+
+            let opaques_unchanged = !self.stalled_on_empty_opaques
+                || infcx.inner.borrow_mut().opaque_types().is_empty();
+
+            if sub_roots_unchanged && opaques_unchanged {
+                self.last_stalled_goal_sub_revision = sub_revision;
+                return errors;
+            }
+        }
+
         loop {
+            let (pass_revision, pass_sub_revision) = infcx.stalled_goal_revisions();
+
             let mut any_changed = false;
             let mut overflowed = false;
+            let mut all_pending_trackable = true;
+            let mut stalled_on_empty_opaques = false;
+            let stalled_sub_roots = &mut self.stalled_sub_roots;
 
             self.obligations.pending.retain_mut(|(obligation, opt_stalled_on)| {
                 if overflowed {
@@ -209,6 +304,14 @@ where
                 if let Some(stalled_on) = opt_stalled_on
                     && delegate.goal_remains_stalled(stalled_on)
                 {
+                    if !Self::record_trackable_stalled_on(
+                        stalled_on,
+                        stalled_sub_roots,
+                        &mut stalled_on_empty_opaques,
+                    ) {
+                        all_pending_trackable = false;
+                    }
+
                     return true;
                 }
 
@@ -279,17 +382,38 @@ where
                         // Update `opt_stalled_on` goal, for the next retain_mut, because we are
                         // running until a fixpoint.
                         *opt_stalled_on = stalled_on;
+
+                        if let Some(stalled_on) = opt_stalled_on {
+                            if !Self::record_trackable_stalled_on(
+                                stalled_on,
+                                stalled_sub_roots,
+                                &mut stalled_on_empty_opaques,
+                            ) {
+                                all_pending_trackable = false;
+                            }
+                        } else {
+                            all_pending_trackable = false;
+                        }
+
                         true
                     }
                 }
             });
             if overflowed {
+                self.all_pending_trackable = false;
                 self.obligations.on_fulfillment_overflow(infcx);
                 // Only return true errors that we have accumulated while processing.
                 return errors;
             }
 
             if !any_changed {
+                self.all_pending_trackable = all_pending_trackable;
+                self.stalled_on_empty_opaques = stalled_on_empty_opaques;
+                self.last_stalled_goal_revision = pass_revision;
+                self.last_stalled_goal_sub_revision = pass_sub_revision;
+
+                self.stalled_sub_roots.retain(|&vid| infcx.stalled_goal_sub_var_is_root(vid));
+
                 break;
             }
         }