Unify CoroutineWitness sooner in typeck, and stall coroutine obligations based off of TypingEnv

compiler/rustc_hir_typeck/src/closure.rs

@@ -161,16 +161,7 @@ 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);
 
-                // 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 = if self.next_trait_solver() {
-                    Ty::new_coroutine_witness(tcx, expr_def_id.to_def_id(), parent_args)
-                } else {
-                    self.next_ty_var(expr_span)
-                };
-
-                self.deferred_coroutine_interiors.borrow_mut().push((expr_def_id, interior));
+                let interior = Ty::new_coroutine_witness(tcx, expr_def_id.to_def_id(), parent_args);
 
                 // Coroutines that come from coroutine closures have not yet determined
                 // their kind ty, so make a fresh infer var which will be constrained

compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs

@@ -631,50 +631,23 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         // trigger query cycle ICEs, as doing so requires MIR.
         self.select_obligations_where_possible(|_| {});
 
-        let coroutines = std::mem::take(&mut *self.deferred_coroutine_interiors.borrow_mut());
-        debug!(?coroutines);
-
-        let mut obligations = vec![];
-
-        if !self.next_trait_solver() {
-            for &(coroutine_def_id, interior) in coroutines.iter() {
-                debug!(?coroutine_def_id);
-
-                // Create the `CoroutineWitness` type that we will unify with `interior`.
-                let args = ty::GenericArgs::identity_for_item(
-                    self.tcx,
-                    self.tcx.typeck_root_def_id(coroutine_def_id.to_def_id()),
-                );
-                let witness =
-                    Ty::new_coroutine_witness(self.tcx, coroutine_def_id.to_def_id(), args);
-
-                // Unify `interior` with `witness` and collect all the resulting obligations.
-                let span = self.tcx.hir_body_owned_by(coroutine_def_id).value.span;
-                let ty::Infer(ty::InferTy::TyVar(_)) = interior.kind() else {
-                    span_bug!(span, "coroutine interior witness not infer: {:?}", interior.kind())
-                };
-                let ok = self
-                    .at(&self.misc(span), self.param_env)
-                    // Will never define opaque types, as all we do is instantiate a type variable.
-                    .eq(DefineOpaqueTypes::Yes, interior, witness)
-                    .expect("Failed to unify coroutine interior type");
-
-                obligations.extend(ok.obligations);
-            }
-        }
+        let ty::TypingMode::Analysis { defining_opaque_types_and_generators } = self.typing_mode()
+        else {
+            bug!();
+        };
 
-        if !coroutines.is_empty() {
-            obligations.extend(
+        if defining_opaque_types_and_generators
+            .iter()
+            .any(|def_id| self.tcx.is_coroutine(def_id.to_def_id()))
+        {
+            self.typeck_results.borrow_mut().coroutine_stalled_predicates.extend(
                 self.fulfillment_cx
                     .borrow_mut()
-                    .drain_stalled_obligations_for_coroutines(&self.infcx),
+                    .drain_stalled_obligations_for_coroutines(&self.infcx)
+                    .into_iter()
+                    .map(|o| (o.predicate, o.cause)),
             );
         }
-
-        self.typeck_results
-            .borrow_mut()
-            .coroutine_stalled_predicates
-            .extend(obligations.into_iter().map(|o| (o.predicate, o.cause)));
     }
 
     #[instrument(skip(self), level = "debug")]

compiler/rustc_hir_typeck/src/typeck_root_ctxt.rs

@@ -63,8 +63,6 @@ pub(crate) struct TypeckRootCtxt<'tcx> {
 
     pub(super) deferred_asm_checks: RefCell<Vec<(&'tcx hir::InlineAsm<'tcx>, HirId)>>,
 
-    pub(super) deferred_coroutine_interiors: RefCell<Vec<(LocalDefId, Ty<'tcx>)>>,
-
     pub(super) deferred_repeat_expr_checks:
         RefCell<Vec<(&'tcx hir::Expr<'tcx>, Ty<'tcx>, ty::Const<'tcx>)>>,
 
@@ -103,7 +101,6 @@ impl<'tcx> TypeckRootCtxt<'tcx> {
             deferred_cast_checks: RefCell::new(Vec::new()),
             deferred_transmute_checks: RefCell::new(Vec::new()),
             deferred_asm_checks: RefCell::new(Vec::new()),
-            deferred_coroutine_interiors: RefCell::new(Vec::new()),
             deferred_repeat_expr_checks: RefCell::new(Vec::new()),
             diverging_type_vars: RefCell::new(Default::default()),
             infer_var_info: RefCell::new(Default::default()),

compiler/rustc_middle/src/ty/context.rs

@@ -721,17 +721,13 @@ impl<'tcx> Interner for TyCtxt<'tcx> {
         self,
         defining_anchor: Self::LocalDefId,
     ) -> Self::LocalDefIds {
-        if self.next_trait_solver_globally() {
-            let coroutines_defined_by = self
-                .nested_bodies_within(defining_anchor)
-                .iter()
-                .filter(|def_id| self.is_coroutine(def_id.to_def_id()));
-            self.mk_local_def_ids_from_iter(
-                self.opaque_types_defined_by(defining_anchor).iter().chain(coroutines_defined_by),
-            )
-        } else {
-            self.opaque_types_defined_by(defining_anchor)
-        }
+        let coroutines_defined_by = self
+            .nested_bodies_within(defining_anchor)
+            .iter()
+            .filter(|def_id| self.is_coroutine(def_id.to_def_id()));
+        self.mk_local_def_ids_from_iter(
+            self.opaque_types_defined_by(defining_anchor).iter().chain(coroutines_defined_by),
+        )
     }
 }
 

compiler/rustc_trait_selection/src/infer.rs

@@ -33,8 +33,8 @@ impl<'tcx> InferCtxt<'tcx> {
         let ty = self.resolve_vars_if_possible(ty);
 
         // FIXME(#132279): This should be removed as it causes us to incorrectly
-        // handle opaques in their defining scope.
-        if !self.next_trait_solver() && !(param_env, ty).has_infer() {
+        // handle opaques in their defining scope, and stalled coroutines.
+        if !self.next_trait_solver() && !(param_env, ty).has_infer() && !ty.has_coroutines() {
             return self.tcx.type_is_copy_modulo_regions(self.typing_env(param_env), ty);
         }
 

compiler/rustc_trait_selection/src/solve.rs

@@ -7,7 +7,7 @@ mod normalize;
 mod select;
 
 pub(crate) use delegate::SolverDelegate;
-pub use fulfill::{FulfillmentCtxt, NextSolverError};
+pub use fulfill::{FulfillmentCtxt, NextSolverError, StalledOnCoroutines};
 pub(crate) use normalize::deeply_normalize_for_diagnostics;
 pub use normalize::{
     deeply_normalize, deeply_normalize_with_skipped_universes,

compiler/rustc_trait_selection/src/solve/fulfill.rs

@@ -10,7 +10,8 @@ use rustc_infer::traits::{
     FromSolverError, PredicateObligation, PredicateObligations, TraitEngine,
 };
 use rustc_middle::ty::{
-    self, DelayedSet, Ty, TyCtxt, TypeSuperVisitable, TypeVisitable, TypeVisitor, TypingMode,
+    self, DelayedSet, Ty, TyCtxt, TypeSuperVisitable, TypeVisitable, TypeVisitableExt, TypeVisitor,
+    TypingMode,
 };
 use rustc_next_trait_solver::delegate::SolverDelegate as _;
 use rustc_next_trait_solver::solve::{
@@ -254,7 +255,7 @@ where
         &mut self,
         infcx: &InferCtxt<'tcx>,
     ) -> PredicateObligations<'tcx> {
-        let stalled_generators = match infcx.typing_mode() {
+        let stalled_coroutines = match infcx.typing_mode() {
             TypingMode::Analysis { defining_opaque_types_and_generators } => {
                 defining_opaque_types_and_generators
             }
@@ -264,7 +265,7 @@ where
             | TypingMode::PostAnalysis => return Default::default(),
         };
 
-        if stalled_generators.is_empty() {
+        if stalled_coroutines.is_empty() {
             return Default::default();
         }
 
@@ -275,7 +276,7 @@ where
                         .visit_proof_tree(
                             obl.as_goal(),
                             &mut StalledOnCoroutines {
-                                stalled_generators,
+                                stalled_coroutines,
                                 span: obl.cause.span,
                                 cache: Default::default(),
                             },
@@ -297,10 +298,10 @@ where
 ///
 /// This function can be also return false positives, which will lead to poor diagnostics
 /// so we want to keep this visitor *precise* too.
-struct StalledOnCoroutines<'tcx> {
-    stalled_generators: &'tcx ty::List<LocalDefId>,
-    span: Span,
-    cache: DelayedSet<Ty<'tcx>>,
+pub struct StalledOnCoroutines<'tcx> {
+    pub stalled_coroutines: &'tcx ty::List<LocalDefId>,
+    pub span: Span,
+    pub cache: DelayedSet<Ty<'tcx>>,
 }
 
 impl<'tcx> inspect::ProofTreeVisitor<'tcx> for StalledOnCoroutines<'tcx> {
@@ -330,12 +331,14 @@ impl<'tcx> TypeVisitor<TyCtxt<'tcx>> for StalledOnCoroutines<'tcx> {
         }
 
         if let ty::CoroutineWitness(def_id, _) = *ty.kind()
-            && def_id.as_local().is_some_and(|def_id| self.stalled_generators.contains(&def_id))
+            && def_id.as_local().is_some_and(|def_id| self.stalled_coroutines.contains(&def_id))
         {
-            return ControlFlow::Break(());
+            ControlFlow::Break(())
+        } else if ty.has_coroutines() {
+            ty.super_visit_with(self)
+        } else {
+            ControlFlow::Continue(())
         }
-
-        ty.super_visit_with(self)
     }
 }
 

compiler/rustc_trait_selection/src/traits/fulfill.rs

@@ -3,6 +3,7 @@ use std::marker::PhantomData;
 use rustc_data_structures::obligation_forest::{
     Error, ForestObligation, ObligationForest, ObligationProcessor, Outcome, ProcessResult,
 };
+use rustc_hir::def_id::LocalDefId;
 use rustc_infer::infer::DefineOpaqueTypes;
 use rustc_infer::traits::{
     FromSolverError, PolyTraitObligation, PredicateObligations, ProjectionCacheKey, SelectionError,
@@ -11,7 +12,10 @@ use rustc_infer::traits::{
 use rustc_middle::bug;
 use rustc_middle::ty::abstract_const::NotConstEvaluatable;
 use rustc_middle::ty::error::{ExpectedFound, TypeError};
-use rustc_middle::ty::{self, Binder, Const, GenericArgsRef, TypeVisitableExt, TypingMode};
+use rustc_middle::ty::{
+    self, Binder, Const, GenericArgsRef, TypeVisitable, TypeVisitableExt, TypingMode,
+};
+use rustc_span::DUMMY_SP;
 use thin_vec::ThinVec;
 use tracing::{debug, debug_span, instrument};
 
@@ -24,6 +28,7 @@ use super::{
 };
 use crate::error_reporting::InferCtxtErrorExt;
 use crate::infer::{InferCtxt, TyOrConstInferVar};
+use crate::solve::StalledOnCoroutines;
 use crate::traits::normalize::normalize_with_depth_to;
 use crate::traits::project::{PolyProjectionObligation, ProjectionCacheKeyExt as _};
 use crate::traits::query::evaluate_obligation::InferCtxtExt;
@@ -166,15 +171,33 @@ where
         &mut self,
         infcx: &InferCtxt<'tcx>,
     ) -> PredicateObligations<'tcx> {
-        let mut processor =
-            DrainProcessor { removed_predicates: PredicateObligations::new(), infcx };
+        let stalled_coroutines = match infcx.typing_mode() {
+            TypingMode::Analysis { defining_opaque_types_and_generators } => {
+                defining_opaque_types_and_generators
+            }
+            TypingMode::Coherence
+            | TypingMode::Borrowck { defining_opaque_types: _ }
+            | TypingMode::PostBorrowckAnalysis { defined_opaque_types: _ }
+            | TypingMode::PostAnalysis => return Default::default(),
+        };
+
+        if stalled_coroutines.is_empty() {
+            return Default::default();
+        }
+
+        let mut processor = DrainProcessor {
+            infcx,
+            removed_predicates: PredicateObligations::new(),
+            stalled_coroutines,
+        };
         let outcome: Outcome<_, _> = self.predicates.process_obligations(&mut processor);
         assert!(outcome.errors.is_empty());
         return processor.removed_predicates;
 
         struct DrainProcessor<'a, 'tcx> {
             infcx: &'a InferCtxt<'tcx>,
             removed_predicates: PredicateObligations<'tcx>,
+            stalled_coroutines: &'tcx ty::List<LocalDefId>,
         }
 
         impl<'tcx> ObligationProcessor for DrainProcessor<'_, 'tcx> {
@@ -183,10 +206,14 @@ where
             type OUT = Outcome<Self::Obligation, Self::Error>;
 
             fn needs_process_obligation(&self, pending_obligation: &Self::Obligation) -> bool {
-                pending_obligation
-                    .stalled_on
-                    .iter()
-                    .any(|&var| self.infcx.ty_or_const_infer_var_changed(var))
+                self.infcx
+                    .resolve_vars_if_possible(pending_obligation.obligation.predicate)
+                    .visit_with(&mut StalledOnCoroutines {
+                        stalled_coroutines: self.stalled_coroutines,
+                        span: DUMMY_SP,
+                        cache: Default::default(),
+                    })
+                    .is_break()
             }
 
             fn process_obligation(

compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs

@@ -847,6 +847,14 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                     }
                 }
 
+                ty::CoroutineWitness(def_id, _) => {
+                    if self.should_stall_coroutine_witness(def_id) {
+                        candidates.ambiguous = true;
+                    } else {
+                        candidates.vec.push(AutoImplCandidate);
+                    }
+                }
+
                 ty::Bool
                 | ty::Char
                 | ty::Int(_)
@@ -866,7 +874,6 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                 | ty::Coroutine(..)
                 | ty::Never
                 | ty::Tuple(_)
-                | ty::CoroutineWitness(..)
                 | ty::UnsafeBinder(_) => {
                     // Only consider auto impls of unsafe traits when there are
                     // no unsafe fields.

compiler/rustc_trait_selection/src/traits/select/mod.rs

@@ -1508,7 +1508,8 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                 defining_opaque_types_and_generators: defining_opaque_types,
             }
             | TypingMode::Borrowck { defining_opaque_types } => {
-                defining_opaque_types.is_empty() || !pred.has_opaque_types()
+                defining_opaque_types.is_empty()
+                    || (!pred.has_opaque_types() && !pred.has_coroutines())
             }
             // The hidden types of `defined_opaque_types` is not local to the current
             // inference context, so we can freely move this to the global cache.
@@ -2239,13 +2240,17 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
             }
 
             ty::CoroutineWitness(def_id, args) => {
-                let hidden_types = rebind_coroutine_witness_types(
-                    self.infcx.tcx,
-                    def_id,
-                    args,
-                    obligation.predicate.bound_vars(),
-                );
-                Where(hidden_types)
+                if self.should_stall_coroutine_witness(def_id) {
+                    Ambiguous
+                } else {
+                    let hidden_types = rebind_coroutine_witness_types(
+                        self.infcx.tcx,
+                        def_id,
+                        args,
+                        obligation.predicate.bound_vars(),
+                    );
+                    Where(hidden_types)
+                }
             }
 
             ty::Closure(_, args) => {
@@ -2871,6 +2876,22 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
 
         obligations
     }
+
+    fn should_stall_coroutine_witness(&self, def_id: DefId) -> bool {
+        match self.infcx.typing_mode() {
+            TypingMode::Analysis { defining_opaque_types_and_generators: stalled_generators } => {
+                if def_id.as_local().is_some_and(|def_id| stalled_generators.contains(&def_id)) {
+                    return true;
+                }
+            }
+            TypingMode::Coherence
+            | TypingMode::PostAnalysis
+            | TypingMode::Borrowck { defining_opaque_types: _ }
+            | TypingMode::PostBorrowckAnalysis { defined_opaque_types: _ } => {}
+        }
+
+        false
+    }
 }
 
 fn rebind_coroutine_witness_types<'tcx>(