also filter obligations by sub_root

Author: lcnr

compiler/rustc_hir_typeck/src/fn_ctxt/inspect_obligations.rs

@@ -76,7 +76,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(sub_root_var);
         debug!(?obligations);
         let mut obligations_for_self_ty = PredicateObligations::new();
         for obligation in obligations {

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;
@@ -106,6 +106,15 @@ pub trait TraitEngine<'tcx, E: 'tcx>: 'tcx {
     fn has_pending_obligations(&self) -> bool;
 
     fn pending_obligations(&self) -> PredicateObligations<'tcx>;
+    // Returning all pending obligations which reference an inference
+    // variable with `_sub_root`. This assumes that no type inference
+    // progress has been made since the last `select_where_possible` call.
+    fn pending_obligations_potentially_referencing_sub_root(
+        &self,
+        _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 `select_where_possible`. Those obligations are marked as

compiler/rustc_trait_selection/src/solve/fulfill.rs

@@ -10,8 +10,8 @@ use rustc_infer::traits::{
     FromSolverError, PredicateObligation, PredicateObligations, TraitEngine,
 };
 use rustc_middle::ty::{
-    self, DelayedSet, Ty, TyCtxt, TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeVisitor,
-    TypingMode,
+    self, DelayedSet, Ty, TyCtxt, TyVid, TypeSuperVisitable, TypeVisitable, TypeVisitableExt,
+    TypeVisitor, TypingMode,
 };
 use rustc_next_trait_solver::delegate::SolverDelegate as _;
 use rustc_next_trait_solver::solve::{
@@ -85,8 +85,30 @@ impl<'tcx> ObligationStorage<'tcx> {
         obligations.extend(self.overflowed.iter().cloned());
         obligations
     }
+    fn clone_pending_potentially_referencing_sub_root(
+        &self,
+        vid: TyVid,
+    ) -> PredicateObligations<'tcx> {
+        let mut obligations: PredicateObligations<'tcx> = self
+            .pending
+            .iter()
+            .filter(|(_, stalled_on)| {
+                // This may incorrectly return `false` in case we've made
+                // inference progress since the last time we tried to evaluate
+                // this obligation.
+                if let Some(stalled_on) = stalled_on {
+                    stalled_on.sub_roots.iter().any(|&r| r == vid)
+                } else {
+                    true
+                }
+            })
+            .map(|(o, _)| o.clone())
+            .collect();
+        obligations.extend(self.overflowed.iter().cloned());
+        obligations
+    }
 
-    fn drain_pending(
+    fn drain_pending_ignoring_overflowed(
         &mut self,
         cond: impl Fn(&PredicateObligation<'tcx>) -> bool,
     ) -> PendingObligations<'tcx> {
@@ -184,7 +206,9 @@ where
         let mut errors = Vec::new();
         loop {
             let mut any_changed = false;
-            for (mut obligation, stalled_on) in self.obligations.drain_pending(|_| true) {
+            for (mut obligation, stalled_on) in
+                self.obligations.drain_pending_ignoring_overflowed(|_| true)
+            {
                 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.
@@ -277,6 +301,12 @@ where
     fn pending_obligations(&self) -> PredicateObligations<'tcx> {
         self.obligations.clone_pending()
     }
+    fn pending_obligations_potentially_referencing_sub_root(
+        &self,
+        vid: ty::TyVid,
+    ) -> PredicateObligations<'tcx> {
+        self.obligations.clone_pending_potentially_referencing_sub_root(vid)
+    }
 
     fn drain_stalled_obligations_for_coroutines(
         &mut self,
@@ -297,7 +327,7 @@ where
         }
 
         self.obligations
-            .drain_pending(|obl| {
+            .drain_pending_ignoring_overflowed(|obl| {
                 infcx.probe(|_| {
                     infcx
                         .visit_proof_tree(