More

3 files changed, 33 insertions, 34 deletions

diff --git a/compiler/rustc_next_trait_solver/src/solve/mod.rs b/compiler/rustc_next_trait_solver/src/solve/mod.rs
index 8868769906d..0695c5acdca 100644
--- a/compiler/rustc_next_trait_solver/src/solve/mod.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/mod.rs
@@ -329,10 +329,7 @@ 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,
-                stalled_generators: _,
-            }
+            TypingMode::Analysis { defining_opaque_types_and_generators: non_rigid_opaques }
             | 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))
diff --git a/compiler/rustc_next_trait_solver/src/solve/normalizes_to/opaque_types.rs b/compiler/rustc_next_trait_solver/src/solve/normalizes_to/opaque_types.rs
index dcfc3b3a701..ee439f1b3d0 100644
--- a/compiler/rustc_next_trait_solver/src/solve/normalizes_to/opaque_types.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/normalizes_to/opaque_types.rs
@@ -33,11 +33,11 @@ where
                 );
                 self.evaluate_added_goals_and_make_canonical_response(Certainty::AMBIGUOUS)
             }
-            TypingMode::Analysis { defining_opaque_types, stalled_generators: _ } => {
+            TypingMode::Analysis { defining_opaque_types_and_generators } => {
                 let Some(def_id) = opaque_ty
                     .def_id
                     .as_local()
-                    .filter(|&def_id| defining_opaque_types.contains(&def_id))
+                    .filter(|&def_id| defining_opaque_types_and_generators.contains(&def_id))
                 else {
                     self.structurally_instantiate_normalizes_to_term(goal, goal.predicate.alias);
                     return self.evaluate_added_goals_and_make_canonical_response(Certainty::Yes);
diff --git a/compiler/rustc_next_trait_solver/src/solve/trait_goals.rs b/compiler/rustc_next_trait_solver/src/solve/trait_goals.rs
index 13132d60dfa..827853be280 100644
--- a/compiler/rustc_next_trait_solver/src/solve/trait_goals.rs
+++ b/compiler/rustc_next_trait_solver/src/solve/trait_goals.rs
@@ -208,20 +208,9 @@ where
             }
         }
 
-        // 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: _ } => {}
-            }
+        // We need to make sure to stall any coroutines we are inferring to avoid query cycles.
+        if let Some(cand) = ecx.try_stall_coroutine_witness(goal.predicate.self_ty()) {
+            return cand;
         }
 
         ecx.probe_and_evaluate_goal_for_constituent_tys(
@@ -275,20 +264,9 @@ where
             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: _ } => {}
-            }
+        // We need to make sure to stall any coroutines we are inferring to avoid query cycles.
+        if let Some(cand) = ecx.try_stall_coroutine_witness(goal.predicate.self_ty()) {
+            return cand;
         }
 
         ecx.probe_and_evaluate_goal_for_constituent_tys(
@@ -1400,4 +1378,28 @@ where
         let candidates = self.assemble_and_evaluate_candidates(goal, AssembleCandidatesFrom::All);
         self.merge_trait_candidates(goal, candidates)
     }
+
+    fn try_stall_coroutine_witness(
+        &mut self,
+        self_ty: I::Ty,
+    ) -> Option<Result<Candidate<I>, NoSolution>> {
+        if let ty::CoroutineWitness(def_id, _) = self_ty.kind() {
+            match self.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 Some(self.forced_ambiguity(MaybeCause::Ambiguity));
+                    }
+                }
+                TypingMode::Coherence
+                | TypingMode::PostAnalysis
+                | TypingMode::Borrowck { defining_opaque_types: _ }
+                | TypingMode::PostBorrowckAnalysis { defined_opaque_types: _ } => {}
+            }
+        }
+
+        None
+    }
 }