Properly stall coroutine witnesses in new solver

compiler/rustc_hir_typeck/src/closure.rs

@@ -163,7 +163,15 @@ 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);
 
-                let interior = self.next_ty_var(expr_span);
+                // 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));
 
                 // Coroutines that come from coroutine closures have not yet determined

compiler/rustc_hir_typeck/src/fn_ctxt/_impl.rs

@@ -635,34 +635,39 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
 
         let mut obligations = vec![];
 
-        for &(coroutine_def_id, interior) in coroutines.iter() {
-            debug!(?coroutine_def_id);
+        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);
 
-            // 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);
+                // 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);
+            }
         }
 
-        // 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_hir_typeck/src/writeback.rs

@@ -8,6 +8,7 @@ use rustc_data_structures::unord::ExtendUnord;
 use rustc_errors::ErrorGuaranteed;
 use rustc_hir::intravisit::{self, InferKind, Visitor};
 use rustc_hir::{self as hir, AmbigArg, HirId};
+use rustc_infer::traits::solve::Goal;
 use rustc_middle::span_bug;
 use rustc_middle::traits::ObligationCause;
 use rustc_middle::ty::adjustment::{Adjust, Adjustment, PointerCoercion};
@@ -763,7 +764,32 @@ impl<'cx, 'tcx> WritebackCx<'cx, 'tcx> {
         T: TypeFoldable<TyCtxt<'tcx>>,
     {
         let value = self.fcx.resolve_vars_if_possible(value);
-        let value = value.fold_with(&mut Resolver::new(self.fcx, span, self.body, true));
+
+        let mut goals = vec![];
+        let value =
+            value.fold_with(&mut Resolver::new(self.fcx, span, self.body, true, &mut goals));
+
+        // Ensure that we resolve goals we get from normalizing coroutine interiors,
+        // but we shouldn't expect those goals to need normalizing (or else we'd get
+        // into a somewhat awkward fixpoint situation, and we don't need it anyways).
+        let mut unexpected_goals = vec![];
+        self.typeck_results.coroutine_stalled_predicates.extend(
+            goals
+                .into_iter()
+                .map(|pred| {
+                    self.fcx.resolve_vars_if_possible(pred).fold_with(&mut Resolver::new(
+                        self.fcx,
+                        span,
+                        self.body,
+                        false,
+                        &mut unexpected_goals,
+                    ))
+                })
+                // FIXME: throwing away the param-env :(
+                .map(|goal| (goal.predicate, self.fcx.misc(span.to_span(self.fcx.tcx)))),
+        );
+        assert_eq!(unexpected_goals, vec![]);
+
         assert!(!value.has_infer());
 
         // We may have introduced e.g. `ty::Error`, if inference failed, make sure
@@ -781,7 +807,12 @@ impl<'cx, 'tcx> WritebackCx<'cx, 'tcx> {
         T: TypeFoldable<TyCtxt<'tcx>>,
     {
         let value = self.fcx.resolve_vars_if_possible(value);
-        let value = value.fold_with(&mut Resolver::new(self.fcx, span, self.body, false));
+
+        let mut goals = vec![];
+        let value =
+            value.fold_with(&mut Resolver::new(self.fcx, span, self.body, false, &mut goals));
+        assert_eq!(goals, vec![]);
+
         assert!(!value.has_infer());
 
         // We may have introduced e.g. `ty::Error`, if inference failed, make sure
@@ -818,6 +849,7 @@ struct Resolver<'cx, 'tcx> {
     /// Whether we should normalize using the new solver, disabled
     /// both when using the old solver and when resolving predicates.
     should_normalize: bool,
+    nested_goals: &'cx mut Vec<Goal<'tcx, ty::Predicate<'tcx>>>,
 }
 
 impl<'cx, 'tcx> Resolver<'cx, 'tcx> {
@@ -826,8 +858,9 @@ impl<'cx, 'tcx> Resolver<'cx, 'tcx> {
         span: &'cx dyn Locatable,
         body: &'tcx hir::Body<'tcx>,
         should_normalize: bool,
+        nested_goals: &'cx mut Vec<Goal<'tcx, ty::Predicate<'tcx>>>,
     ) -> Resolver<'cx, 'tcx> {
-        Resolver { fcx, span, body, should_normalize }
+        Resolver { fcx, span, body, nested_goals, should_normalize }
     }
 
     fn report_error(&self, p: impl Into<ty::GenericArg<'tcx>>) -> ErrorGuaranteed {
@@ -864,12 +897,18 @@ impl<'cx, 'tcx> Resolver<'cx, 'tcx> {
             let cause = ObligationCause::misc(self.span.to_span(tcx), body_id);
             let at = self.fcx.at(&cause, self.fcx.param_env);
             let universes = vec![None; outer_exclusive_binder(value).as_usize()];
-            solve::deeply_normalize_with_skipped_universes(at, value, universes).unwrap_or_else(
-                |errors| {
+            match solve::deeply_normalize_with_skipped_universes_and_ambiguous_goals(
+                at, value, universes,
+            ) {
+                Ok((value, goals)) => {
+                    self.nested_goals.extend(goals);
+                    value
+                }
+                Err(errors) => {
                     let guar = self.fcx.err_ctxt().report_fulfillment_errors(errors);
                     new_err(tcx, guar)
-                },
-            )
+                }
+            }
         } else {
             value
         };

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,17 @@ 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 {
+        if self.next_trait_solver_globally() {
+            self.stalled_generators_within(defining_anchor)
+        } else {
+            ty::List::empty()
+        }
+    }
 }
 
 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,22 @@
             }
         }
 
+        // TODO:
+        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,
@@ -259,6 +275,22 @@
             return Err(NoSolution);
         }
 
+        // TODO:
+        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/infer.rs

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

compiler/rustc_trait_selection/src/solve.rs

@@ -9,5 +9,8 @@ mod select;
 pub(crate) use delegate::SolverDelegate;
 pub use fulfill::{FulfillmentCtxt, NextSolverError};
 pub(crate) use normalize::deeply_normalize_for_diagnostics;
-pub use normalize::{deeply_normalize, deeply_normalize_with_skipped_universes};
+pub use normalize::{
+    deeply_normalize, deeply_normalize_with_skipped_universes,
+    deeply_normalize_with_skipped_universes_and_ambiguous_goals,
+};
 pub use select::InferCtxtSelectExt;