Properly drain pending obligations for coroutines

5 files changed, 122 insertions, 38 deletions

diff --git a/compiler/rustc_trait_selection/src/solve/fulfill.rs b/compiler/rustc_trait_selection/src/solve/fulfill.rs
index 192e632a2d5..abee5ac52c1 100644
--- a/compiler/rustc_trait_selection/src/solve/fulfill.rs
+++ b/compiler/rustc_trait_selection/src/solve/fulfill.rs
@@ -1,18 +1,25 @@
 use std::marker::PhantomData;
 use std::mem;
+use std::ops::ControlFlow;
 
 use rustc_data_structures::thinvec::ExtractIf;
+use rustc_hir::def_id::LocalDefId;
 use rustc_infer::infer::InferCtxt;
 use rustc_infer::traits::query::NoSolution;
 use rustc_infer::traits::{
     FromSolverError, PredicateObligation, PredicateObligations, TraitEngine,
 };
+use rustc_middle::ty::{
+    self, Ty, TyCtxt, TypeSuperVisitable, TypeVisitable, TypeVisitor, TypingMode,
+};
 use rustc_next_trait_solver::solve::{GenerateProofTree, HasChanged, SolverDelegateEvalExt as _};
+use rustc_span::Span;
 use tracing::instrument;
 
 use self::derive_errors::*;
 use super::Certainty;
 use super::delegate::SolverDelegate;
+use super::inspect::{self, ProofTreeInferCtxtExt};
 use crate::traits::{FulfillmentError, ScrubbedTraitError};
 
 mod derive_errors;
@@ -39,7 +46,7 @@ pub struct FulfillmentCtxt<'tcx, E: 'tcx> {
     _errors: PhantomData<E>,
 }
 
-#[derive(Default)]
+#[derive(Default, Debug)]
 struct ObligationStorage<'tcx> {
     /// Obligations which resulted in an overflow in fulfillment itself.
     ///
@@ -55,20 +62,23 @@ impl<'tcx> ObligationStorage<'tcx> {
         self.pending.push(obligation);
     }
 
+    fn has_pending_obligations(&self) -> bool {
+        !self.pending.is_empty() || !self.overflowed.is_empty()
+    }
+
     fn clone_pending(&self) -> PredicateObligations<'tcx> {
         let mut obligations = self.pending.clone();
         obligations.extend(self.overflowed.iter().cloned());
         obligations
     }
 
-    fn take_pending(&mut self) -> PredicateObligations<'tcx> {
-        let mut obligations = mem::take(&mut self.pending);
-        obligations.append(&mut self.overflowed);
-        obligations
-    }
-
-    fn unstalled_for_select(&mut self) -> impl Iterator<Item = PredicateObligation<'tcx>> + 'tcx {
-        mem::take(&mut self.pending).into_iter()
+    fn drain_pending(
+        &mut self,
+        cond: impl Fn(&PredicateObligation<'tcx>) -> bool,
+    ) -> PredicateObligations<'tcx> {
+        let (unstalled, pending) = mem::take(&mut self.pending).into_iter().partition(cond);
+        self.pending = pending;
+        unstalled
     }
 
     fn on_fulfillment_overflow(&mut self, infcx: &InferCtxt<'tcx>) {
@@ -160,7 +170,7 @@ where
             }
 
             let mut has_changed = false;
-            for obligation in self.obligations.unstalled_for_select() {
+            for obligation in self.obligations.drain_pending(|_| true) {
                 let goal = obligation.as_goal();
                 let result = <&SolverDelegate<'tcx>>::from(infcx)
                     .evaluate_root_goal(goal, GenerateProofTree::No, obligation.cause.span)
@@ -196,15 +206,79 @@ where
     }
 
     fn has_pending_obligations(&self) -> bool {
-        !self.obligations.pending.is_empty() || !self.obligations.overflowed.is_empty()
+        self.obligations.has_pending_obligations()
     }
 
     fn pending_obligations(&self) -> PredicateObligations<'tcx> {
         self.obligations.clone_pending()
     }
 
-    fn drain_unstalled_obligations(&mut self, _: &InferCtxt<'tcx>) -> PredicateObligations<'tcx> {
-        self.obligations.take_pending()
+    fn drain_stalled_obligations_for_coroutines(
+        &mut self,
+        infcx: &InferCtxt<'tcx>,
+    ) -> PredicateObligations<'tcx> {
+        self.obligations.drain_pending(|obl| {
+            let stalled_generators = match infcx.typing_mode() {
+                TypingMode::Analysis { defining_opaque_types: _, stalled_generators } => {
+                    stalled_generators
+                }
+                TypingMode::Coherence
+                | TypingMode::Borrowck { defining_opaque_types: _ }
+                | TypingMode::PostBorrowckAnalysis { defined_opaque_types: _ }
+                | TypingMode::PostAnalysis => return false,
+            };
+
+            if stalled_generators.is_empty() {
+                return false;
+            }
+
+            infcx.probe(|_| {
+                infcx
+                    .visit_proof_tree(
+                        obl.as_goal(),
+                        &mut StalledOnCoroutines { stalled_generators, span: obl.cause.span },
+                    )
+                    .is_break()
+            })
+        })
+    }
+}
+
+struct StalledOnCoroutines<'tcx> {
+    stalled_generators: &'tcx ty::List<LocalDefId>,
+    span: Span,
+    // TODO: Cache
+}
+
+impl<'tcx> inspect::ProofTreeVisitor<'tcx> for StalledOnCoroutines<'tcx> {
+    type Result = ControlFlow<()>;
+
+    fn span(&self) -> rustc_span::Span {
+        self.span
+    }
+
+    fn visit_goal(&mut self, inspect_goal: &super::inspect::InspectGoal<'_, 'tcx>) -> Self::Result {
+        inspect_goal.goal().predicate.visit_with(self)?;
+
+        if let Some(candidate) = inspect_goal.unique_applicable_candidate() {
+            candidate.visit_nested_no_probe(self)
+        } else {
+            ControlFlow::Continue(())
+        }
+    }
+}
+
+impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for StalledOnCoroutines<'tcx> {
+    type Result = ControlFlow<()>;
+
+    fn visit_ty(&mut self, ty: Ty<'tcx>) -> Self::Result {
+        if let ty::CoroutineWitness(def_id, _) = *ty.kind()
+            && def_id.as_local().is_some_and(|def_id| self.stalled_generators.contains(&def_id))
+        {
+            return ControlFlow::Break(());
+        }
+
+        ty.super_visit_with(self)
     }
 }
 
diff --git a/compiler/rustc_trait_selection/src/solve/fulfill/derive_errors.rs b/compiler/rustc_trait_selection/src/solve/fulfill/derive_errors.rs
index d8dcd12aecb..e31cad6b40d 100644
--- a/compiler/rustc_trait_selection/src/solve/fulfill/derive_errors.rs
+++ b/compiler/rustc_trait_selection/src/solve/fulfill/derive_errors.rs
@@ -109,10 +109,16 @@ pub(super) fn fulfillment_error_for_stalled<'tcx>(
                 false,
             ),
             Ok((_, Certainty::Yes)) => {
-                bug!("did not expect successful goal when collecting ambiguity errors")
+                bug!(
+                    "did not expect successful goal when collecting ambiguity errors for `{:?}`",
+                    infcx.resolve_vars_if_possible(root_obligation.predicate),
+                )
             }
             Err(_) => {
-                bug!("did not expect selection error when collecting ambiguity errors")
+                bug!(
+                    "did not expect selection error when collecting ambiguity errors for `{:?}`",
+                    infcx.resolve_vars_if_possible(root_obligation.predicate),
+                )
             }
         }
     });
diff --git a/compiler/rustc_trait_selection/src/solve/normalize.rs b/compiler/rustc_trait_selection/src/solve/normalize.rs
index 7b80059722a..65ab14ae07c 100644
--- a/compiler/rustc_trait_selection/src/solve/normalize.rs
+++ b/compiler/rustc_trait_selection/src/solve/normalize.rs
@@ -76,7 +76,7 @@ where
     let value = value.try_fold_with(&mut folder)?;
     let goals = folder
         .fulfill_cx
-        .drain_unstalled_obligations(at.infcx)
+        .drain_stalled_obligations_for_coroutines(at.infcx)
         .into_iter()
         .map(|obl| obl.as_goal())
         .collect();
@@ -130,7 +130,7 @@ where
         );
 
         self.fulfill_cx.register_predicate_obligation(infcx, obligation);
-        let errors = self.fulfill_cx.select_all_or_error(infcx);
+        let errors = self.fulfill_cx.select_where_possible(infcx);
         if !errors.is_empty() {
             return Err(errors);
         }
@@ -171,7 +171,7 @@ where
 
         let result = if infcx.predicate_may_hold(&obligation) {
             self.fulfill_cx.register_predicate_obligation(infcx, obligation);
-            let errors = self.fulfill_cx.select_all_or_error(infcx);
+            let errors = self.fulfill_cx.select_where_possible(infcx);
             if !errors.is_empty() {
                 return Err(errors);
             }
@@ -286,27 +286,31 @@ impl<'tcx> TypeFolder<TyCtxt<'tcx>> for DeeplyNormalizeForDiagnosticsFolder<'_,
 
     fn fold_ty(&mut self, ty: Ty<'tcx>) -> Ty<'tcx> {
         let infcx = self.at.infcx;
-        infcx
-            .commit_if_ok(|_| {
-                deeply_normalize_with_skipped_universes(
-                    self.at,
-                    ty,
-                    vec![None; ty.outer_exclusive_binder().as_usize()],
-                )
-            })
-            .unwrap_or_else(|_: Vec<ScrubbedTraitError<'tcx>>| ty.super_fold_with(self))
+        let result =
+            infcx.commit_if_ok(|_| {
+                deeply_normalize_with_skipped_universes_and_ambiguous_goals::<
+                    _,
+                    ScrubbedTraitError<'tcx>,
+                >(self.at, ty, vec![None; ty.outer_exclusive_binder().as_usize()])
+            });
+        match result {
+            Ok((ty, _)) => ty,
+            Err(_) => ty.super_fold_with(self),
+        }
     }
 
     fn fold_const(&mut self, ct: ty::Const<'tcx>) -> ty::Const<'tcx> {
         let infcx = self.at.infcx;
-        infcx
-            .commit_if_ok(|_| {
-                deeply_normalize_with_skipped_universes(
-                    self.at,
-                    ct,
-                    vec![None; ct.outer_exclusive_binder().as_usize()],
-                )
-            })
-            .unwrap_or_else(|_: Vec<ScrubbedTraitError<'tcx>>| ct.super_fold_with(self))
+        let result =
+            infcx.commit_if_ok(|_| {
+                deeply_normalize_with_skipped_universes_and_ambiguous_goals::<
+                    _,
+                    ScrubbedTraitError<'tcx>,
+                >(self.at, ct, vec![None; ct.outer_exclusive_binder().as_usize()])
+            });
+        match result {
+            Ok((ct, _)) => ct,
+            Err(_) => ct.super_fold_with(self),
+        }
     }
 }
diff --git a/compiler/rustc_trait_selection/src/traits/fulfill.rs b/compiler/rustc_trait_selection/src/traits/fulfill.rs
index 1b76d48e431..a11f8d3a9ec 100644
--- a/compiler/rustc_trait_selection/src/traits/fulfill.rs
+++ b/compiler/rustc_trait_selection/src/traits/fulfill.rs
@@ -162,7 +162,7 @@ where
         self.select(selcx)
     }
 
-    fn drain_unstalled_obligations(
+    fn drain_stalled_obligations_for_coroutines(
         &mut self,
         infcx: &InferCtxt<'tcx>,
     ) -> PredicateObligations<'tcx> {
diff --git a/compiler/rustc_trait_selection/src/traits/select/mod.rs b/compiler/rustc_trait_selection/src/traits/select/mod.rs
index 4d88a23250a..5255b57c791 100644
--- a/compiler/rustc_trait_selection/src/traits/select/mod.rs
+++ b/compiler/rustc_trait_selection/src/traits/select/mod.rs
@@ -1498,7 +1498,7 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
             // However, if we disqualify *all* goals from being cached, perf suffers.
             // This is likely fixed by better caching in general in the new solver.
             // See: <https://github.com/rust-lang/rust/issues/132064>.
-            TypingMode::Analysis { defining_opaque_types }
+            TypingMode::Analysis { defining_opaque_types, stalled_generators: _ }
             | TypingMode::Borrowck { defining_opaque_types } => {
                 defining_opaque_types.is_empty() || !pred.has_opaque_types()
             }