obligations_for_self_ty: skip irrelevant goals (recompute sub_root from stalled_vars)

compiler/rustc_hir_typeck/src/fn_ctxt/inspect_obligations.rs

@@ -38,10 +38,10 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     ) -> bool {
         match predicate.kind().skip_binder() {
             ty::PredicateKind::Clause(ty::ClauseKind::Trait(data)) => {
-                self.type_matches_expected_vid(expected_vid, data.self_ty())
+                self.type_matches_expected_vid(data.self_ty(), expected_vid)
             }
             ty::PredicateKind::Clause(ty::ClauseKind::Projection(data)) => {
-                self.type_matches_expected_vid(expected_vid, data.projection_term.self_ty())
+                self.type_matches_expected_vid(data.projection_term.self_ty(), expected_vid)
             }
             ty::PredicateKind::Clause(ty::ClauseKind::ConstArgHasType(..))
             | ty::PredicateKind::Subtype(..)
@@ -61,7 +61,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     }
 
     #[instrument(level = "debug", skip(self), ret)]
-    fn type_matches_expected_vid(&self, expected_vid: ty::TyVid, ty: Ty<'tcx>) -> bool {
+    fn type_matches_expected_vid(&self, ty: Ty<'tcx>, expected_vid: ty::TyVid) -> bool {
         let ty = self.shallow_resolve(ty);
         debug!(?ty);
 
@@ -77,7 +77,18 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         &self,
         self_ty: ty::TyVid,
     ) -> PredicateObligations<'tcx> {
-        let obligations = self.fulfillment_cx.borrow().pending_obligations();
+        // We only look at obligations which may reference the self type.
+        // This lookup uses the `sub_root` instead of the inference variable
+        // itself as that's slightly nicer to implement. It shouldn't really
+        // matter.
+        //
+        // This is really impactful when typechecking functions with a lot of
+        // stalled obligations, e.g. in the `wg-grammar` benchmark.
+        let sub_root_var = self.sub_unification_table_root_var(self_ty);
+        let obligations = self
+            .fulfillment_cx
+            .borrow()
+            .pending_obligations_potentially_referencing_sub_root(&self.infcx, sub_root_var);
         debug!(?obligations);
         let mut obligations_for_self_ty = PredicateObligations::new();
         for obligation in obligations {
@@ -189,6 +200,18 @@ impl<'tcx> ProofTreeVisitor<'tcx> for NestedObligationsForSelfTy<'_, 'tcx> {
             return;
         }
 
+        // We don't care about any pending goals which don't actually
+        // use the self type.
+        if !inspect_goal
+            .orig_values()
+            .iter()
+            .filter_map(|arg| arg.as_type())
+            .any(|ty| self.fcx.type_matches_expected_vid(ty, self.self_ty))
+        {
+            debug!(goal = ?inspect_goal.goal(), "goal does not mention self type");
+            return;
+        }
+
         let tcx = self.fcx.tcx;
         let goal = inspect_goal.goal();
         if self.fcx.predicate_has_self_ty(goal.predicate, self.self_ty) {

compiler/rustc_infer/src/traits/engine.rs

@@ -1,7 +1,7 @@
 use std::fmt::Debug;
 
 use rustc_hir::def_id::DefId;
-use rustc_middle::ty::{self, Ty, Upcast};
+use rustc_middle::ty::{self, Ty, TyVid, Upcast};
 
 use super::{ObligationCause, PredicateObligation, PredicateObligations};
 use crate::infer::InferCtxt;
@@ -108,6 +108,16 @@ pub trait TraitEngine<'tcx, E: 'tcx>: 'tcx {
 
     fn pending_obligations(&self) -> PredicateObligations<'tcx>;
 
+    /// Returning all pending obligations which reference an inference
+    /// variable with `_sub_root`.
+    fn pending_obligations_potentially_referencing_sub_root(
+        &self,
+        _infcx: &InferCtxt<'tcx>,
+        _sub_root: TyVid,
+    ) -> PredicateObligations<'tcx> {
+        self.pending_obligations()
+    }
+
     /// Among all pending obligations, collect those are stalled on a inference variable which has
     /// changed since the last call to `try_evaluate_obligations`. Those obligations are marked as
     /// successful and returned.

compiler/rustc_trait_selection/src/solve/fulfill.rs

@@ -6,7 +6,7 @@ use rustc_infer::traits::query::NoSolution;
 use rustc_infer::traits::{
     FromSolverError, PredicateObligation, PredicateObligations, TraitEngine,
 };
-use rustc_middle::ty::{TyCtxt, TypeVisitableExt, TypingMode};
+use rustc_middle::ty::{self, TyCtxt, TyVid, TypeVisitableExt, TypingMode};
 use rustc_next_trait_solver::delegate::SolverDelegate as _;
 use rustc_next_trait_solver::solve::{
     GoalEvaluation, GoalStalledOn, HasChanged, MaybeInfo, SolverDelegateEvalExt as _,
@@ -79,6 +79,32 @@ impl<'tcx> ObligationStorage<'tcx> {
         obligations
     }
 
+    fn clone_pending_potentially_referencing_sub_root(
+        &self,
+        infcx: &InferCtxt<'tcx>,
+        vid: TyVid,
+    ) -> PredicateObligations<'tcx> {
+        let mut obligations: PredicateObligations<'tcx> = self
+            .pending
+            .iter()
+            .filter(|(_, stalled_on)| {
+                let Some(stalled_on) = stalled_on else { return true };
+                // Conservative here: if any of its stalled vars are not infer var anymore,
+                // some unification happened, so this goal is no longer stalled.
+                // So include it to re-evaluate in the 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,
+                        _ => true,
+                    },
+                )
+            })
+            .map(|(o, _)| o.clone())
+            .collect();
+        obligations.extend(self.overflowed.iter().cloned());
+        obligations
+    }
+
     fn drain_pending(
         &mut self,
         cond: impl Fn(&PredicateObligation<'tcx>, &Option<GoalStalledOn<TyCtxt<'tcx>>>) -> bool,
@@ -271,6 +297,14 @@ where
         self.obligations.clone_pending()
     }
 
+    fn pending_obligations_potentially_referencing_sub_root(
+        &self,
+        infcx: &InferCtxt<'tcx>,
+        vid: ty::TyVid,
+    ) -> PredicateObligations<'tcx> {
+        self.obligations.clone_pending_potentially_referencing_sub_root(infcx, vid)
+    }
+
     fn drain_stalled_obligations_for_coroutines(
         &mut self,
         infcx: &InferCtxt<'tcx>,

compiler/rustc_trait_selection/src/solve/inspect/analyse.rs

@@ -317,6 +317,10 @@ impl<'a, 'tcx> InspectGoal<'a, 'tcx> {
         self.depth
     }
 
+    pub fn orig_values(&self) -> &[ty::GenericArg<'tcx>] {
+        &self.orig_values
+    }
+
     fn candidates_recur(
         &'a self,
         candidates: &mut Vec<InspectCandidate<'a, 'tcx>>,