Properly drain pending obligations for coroutines

Author: compiler-errors

compiler/rustc_hir_typeck/src/closure.rs

@@ -163,6 +163,9 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 // Resume type defaults to `()` if the coroutine has no argument.
                 let resume_ty = liberated_sig.inputs().get(0).copied().unwrap_or(tcx.types.unit);
 
+                // TODO: In the new solver, we can just instantiate this eagerly
+                // with the witness. This will ensure that goals that don't need
+                // to stall on interior types will get processed eagerly.
                 let interior = self.next_ty_var(expr_span);
                 self.deferred_coroutine_interiors.borrow_mut().push((expr_def_id, interior));
 

compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs

@@ -659,10 +659,12 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             obligations.extend(ok.obligations);
         }
 
-        // FIXME: Use a real visitor for unstalled obligations in the new solver.
         if !coroutines.is_empty() {
-            obligations
-                .extend(self.fulfillment_cx.borrow_mut().drain_unstalled_obligations(&self.infcx));
+            obligations.extend(
+                self.fulfillment_cx
+                    .borrow_mut()
+                    .drain_stalled_obligations_for_coroutines(&self.infcx),
+            );
         }
 
         self.typeck_results

compiler/rustc_hir_typeck/src/typeck_root_ctxt.rs

@@ -84,7 +84,7 @@ impl<'tcx> TypeckRootCtxt<'tcx> {
         let hir_owner = tcx.local_def_id_to_hir_id(def_id).owner;
 
         let infcx =
-            tcx.infer_ctxt().ignoring_regions().build(TypingMode::analysis_in_body(tcx, def_id));
+            tcx.infer_ctxt().ignoring_regions().build(TypingMode::typeck_for_body(tcx, def_id));
         let typeck_results = RefCell::new(ty::TypeckResults::new(hir_owner));
 
         TypeckRootCtxt {

compiler/rustc_infer/src/infer/mod.rs

@@ -967,7 +967,7 @@ impl<'tcx> InferCtxt<'tcx> {
     pub fn can_define_opaque_ty(&self, id: impl Into<DefId>) -> bool {
         debug_assert!(!self.next_trait_solver());
         match self.typing_mode() {
-            TypingMode::Analysis { defining_opaque_types }
+            TypingMode::Analysis { defining_opaque_types, stalled_generators: _ }
             | TypingMode::Borrowck { defining_opaque_types } => {
                 id.into().as_local().is_some_and(|def_id| defining_opaque_types.contains(&def_id))
             }
@@ -1262,7 +1262,7 @@ impl<'tcx> InferCtxt<'tcx> {
             // to handle them without proper canonicalization. This means we may cause cycle
             // errors and fail to reveal opaques while inside of bodies. We should rename this
             // function and require explicit comments on all use-sites in the future.
-            ty::TypingMode::Analysis { defining_opaque_types: _ }
+            ty::TypingMode::Analysis { defining_opaque_types: _, stalled_generators: _ }
             | ty::TypingMode::Borrowck { defining_opaque_types: _ } => {
                 TypingMode::non_body_analysis()
             }

compiler/rustc_infer/src/traits/engine.rs

@@ -94,7 +94,7 @@ pub trait TraitEngine<'tcx, E: 'tcx>: 'tcx {
     /// 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
     /// successful and returned.
-    fn drain_unstalled_obligations(
+    fn drain_stalled_obligations_for_coroutines(
         &mut self,
         infcx: &InferCtxt<'tcx>,
     ) -> PredicateObligations<'tcx>;

compiler/rustc_middle/src/query/mod.rs

@@ -387,6 +387,15 @@ rustc_queries! {
         }
     }
 
+    query stalled_generators_within(
+        key: LocalDefId
+    ) -> &'tcx ty::List<LocalDefId> {
+        desc {
+            |tcx| "computing the opaque types defined by `{}`",
+            tcx.def_path_str(key.to_def_id())
+        }
+    }
+
     /// Returns the explicitly user-written *bounds* on the associated or opaque type given by `DefId`
     /// that must be proven true at definition site (and which can be assumed at usage sites).
     ///

compiler/rustc_middle/src/query/plumbing.rs

@@ -366,11 +366,11 @@ macro_rules! define_callbacks {
 
                 pub type Storage<'tcx> = <$($K)* as keys::Key>::Cache<Erase<$V>>;
 
-                // Ensure that keys grow no larger than 80 bytes by accident.
+                // Ensure that keys grow no larger than 96 bytes by accident.
                 // Increase this limit if necessary, but do try to keep the size low if possible
                 #[cfg(target_pointer_width = "64")]
                 const _: () = {
-                    if size_of::<Key<'static>>() > 88 {
+                    if size_of::<Key<'static>>() > 96 {
                         panic!("{}", concat!(
                             "the query `",
                             stringify!($name),

compiler/rustc_middle/src/ty/context.rs

@@ -106,7 +106,7 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
     ) -> Self::PredefinedOpaques {
         self.mk_predefined_opaques_in_body(data)
     }
-    type DefiningOpaqueTypes = &'tcx ty::List<LocalDefId>;
+    type LocalDefIds = &'tcx ty::List<LocalDefId>;
     type CanonicalVars = CanonicalVarInfos<'tcx>;
     fn mk_canonical_var_infos(self, infos: &[ty::CanonicalVarInfo<Self>]) -> Self::CanonicalVars {
         self.mk_canonical_var_infos(infos)
@@ -674,9 +674,13 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
         self.anonymize_bound_vars(binder)
     }
 
-    fn opaque_types_defined_by(self, defining_anchor: LocalDefId) -> Self::DefiningOpaqueTypes {
+    fn opaque_types_defined_by(self, defining_anchor: LocalDefId) -> Self::LocalDefIds {
         self.opaque_types_defined_by(defining_anchor)
     }
+
+    fn stalled_generators_within(self, defining_anchor: Self::LocalDefId) -> Self::LocalDefIds {
+        self.stalled_generators_within(defining_anchor)
+    }
 }
 
 macro_rules! bidirectional_lang_item_map {

compiler/rustc_next_trait_solver/src/solve/mod.rs

@@ -329,7 +329,10 @@ where
             TypingMode::Coherence | TypingMode::PostAnalysis => false,
             // During analysis, opaques are rigid unless they may be defined by
             // the current body.
-            TypingMode::Analysis { defining_opaque_types: non_rigid_opaques }
+            TypingMode::Analysis {
+                defining_opaque_types: non_rigid_opaques,
+                stalled_generators: _,
+            }
             | TypingMode::Borrowck { defining_opaque_types: non_rigid_opaques }
             | TypingMode::PostBorrowckAnalysis { defined_opaque_types: non_rigid_opaques } => {
                 !def_id.as_local().is_some_and(|def_id| non_rigid_opaques.contains(&def_id))

compiler/rustc_next_trait_solver/src/solve/normalizes_to/opaque_types.rs

@@ -33,7 +33,7 @@ where
                 );
                 self.evaluate_added_goals_and_make_canonical_response(Certainty::AMBIGUOUS)
             }
-            TypingMode::Analysis { defining_opaque_types } => {
+            TypingMode::Analysis { defining_opaque_types, stalled_generators: _ } => {
                 let Some(def_id) = opaque_ty
                     .def_id
                     .as_local()

compiler/rustc_next_trait_solver/src/solve/trait_goals.rs

@@ -208,6 +208,21 @@ where
             }
         }
 
+        if let ty::CoroutineWitness(def_id, _) = goal.predicate.self_ty().kind() {
+            match ecx.typing_mode() {
+                TypingMode::Analysis { stalled_generators, defining_opaque_types: _ } => {
+                    if def_id.as_local().is_some_and(|def_id| stalled_generators.contains(&def_id))
+                    {
+                        return ecx.forced_ambiguity(MaybeCause::Ambiguity);
+                    }
+                }
+                TypingMode::Coherence
+                | TypingMode::PostAnalysis
+                | TypingMode::Borrowck { defining_opaque_types: _ }
+                | TypingMode::PostBorrowckAnalysis { defined_opaque_types: _ } => {}
+            }
+        }
+
         ecx.probe_and_evaluate_goal_for_constituent_tys(
             CandidateSource::BuiltinImpl(BuiltinImplSource::Misc),
             goal,

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)
     }
 }
 

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),
+                )
             }
         }
     });