Rollup merge of #122829 - ShoyuVanilla:gen-block-impl-fused-iter, r=compiler-errors

Implement `FusedIterator` for `gen` block

cc #117078

6 files changed, 80 insertions, 7 deletions

diff --git a/compiler/rustc_trait_selection/src/solve/assembly/mod.rs b/compiler/rustc_trait_selection/src/solve/assembly/mod.rs
index 9f33dce2a6d..d92bae2528f 100644
--- a/compiler/rustc_trait_selection/src/solve/assembly/mod.rs
+++ b/compiler/rustc_trait_selection/src/solve/assembly/mod.rs
@@ -215,6 +215,13 @@ pub(super) trait GoalKind<'tcx>:
         goal: Goal<'tcx, Self>,
     ) -> QueryResult<'tcx>;
 
+    /// A coroutine (that comes from a `gen` desugaring) is known to implement
+    /// `FusedIterator`
+    fn consider_builtin_fused_iterator_candidate(
+        ecx: &mut EvalCtxt<'_, 'tcx>,
+        goal: Goal<'tcx, Self>,
+    ) -> QueryResult<'tcx>;
+
     fn consider_builtin_async_iterator_candidate(
         ecx: &mut EvalCtxt<'_, 'tcx>,
         goal: Goal<'tcx, Self>,
@@ -497,6 +504,8 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             G::consider_builtin_future_candidate(self, goal)
         } else if lang_items.iterator_trait() == Some(trait_def_id) {
             G::consider_builtin_iterator_candidate(self, goal)
+        } else if lang_items.fused_iterator_trait() == Some(trait_def_id) {
+            G::consider_builtin_fused_iterator_candidate(self, goal)
         } else if lang_items.async_iterator_trait() == Some(trait_def_id) {
             G::consider_builtin_async_iterator_candidate(self, goal)
         } else if lang_items.coroutine_trait() == Some(trait_def_id) {
diff --git a/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs b/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs
index 85bb6338daf..66688893235 100644
--- a/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs
+++ b/compiler/rustc_trait_selection/src/solve/normalizes_to/mod.rs
@@ -647,6 +647,13 @@ impl<'tcx> assembly::GoalKind<'tcx> for NormalizesTo<'tcx> {
         )
     }
 
+    fn consider_builtin_fused_iterator_candidate(
+        _ecx: &mut EvalCtxt<'_, 'tcx>,
+        goal: Goal<'tcx, Self>,
+    ) -> QueryResult<'tcx> {
+        bug!("`FusedIterator` does not have an associated type: {:?}", goal);
+    }
+
     fn consider_builtin_async_iterator_candidate(
         ecx: &mut EvalCtxt<'_, 'tcx>,
         goal: Goal<'tcx, Self>,
diff --git a/compiler/rustc_trait_selection/src/solve/trait_goals.rs b/compiler/rustc_trait_selection/src/solve/trait_goals.rs
index c252ad76dfe..184ba31f19d 100644
--- a/compiler/rustc_trait_selection/src/solve/trait_goals.rs
+++ b/compiler/rustc_trait_selection/src/solve/trait_goals.rs
@@ -456,6 +456,28 @@ impl<'tcx> assembly::GoalKind<'tcx> for TraitPredicate<'tcx> {
         ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
     }
 
+    fn consider_builtin_fused_iterator_candidate(
+        ecx: &mut EvalCtxt<'_, 'tcx>,
+        goal: Goal<'tcx, Self>,
+    ) -> QueryResult<'tcx> {
+        if goal.predicate.polarity != ty::ImplPolarity::Positive {
+            return Err(NoSolution);
+        }
+
+        let ty::Coroutine(def_id, _) = *goal.predicate.self_ty().kind() else {
+            return Err(NoSolution);
+        };
+
+        // Coroutines are not iterators unless they come from `gen` desugaring
+        let tcx = ecx.tcx();
+        if !tcx.coroutine_is_gen(def_id) {
+            return Err(NoSolution);
+        }
+
+        // Gen coroutines unconditionally implement `FusedIterator`
+        ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
+    }
+
     fn consider_builtin_async_iterator_candidate(
         ecx: &mut EvalCtxt<'_, 'tcx>,
         goal: Goal<'tcx, Self>,
diff --git a/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs b/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs
index 49091e53be7..9fb4577fb21 100644
--- a/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs
+++ b/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs
@@ -118,6 +118,8 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                     self.assemble_future_candidates(obligation, &mut candidates);
                 } else if lang_items.iterator_trait() == Some(def_id) {
                     self.assemble_iterator_candidates(obligation, &mut candidates);
+                } else if lang_items.fused_iterator_trait() == Some(def_id) {
+                    self.assemble_fused_iterator_candidates(obligation, &mut candidates);
                 } else if lang_items.async_iterator_trait() == Some(def_id) {
                     self.assemble_async_iterator_candidates(obligation, &mut candidates);
                 } else if lang_items.async_fn_kind_helper() == Some(def_id) {
@@ -302,14 +304,31 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
         candidates: &mut SelectionCandidateSet<'tcx>,
     ) {
         let self_ty = obligation.self_ty().skip_binder();
-        if let ty::Coroutine(did, ..) = self_ty.kind() {
-            // gen constructs get lowered to a special kind of coroutine that
-            // should directly `impl Iterator`.
-            if self.tcx().coroutine_is_gen(*did) {
-                debug!(?self_ty, ?obligation, "assemble_iterator_candidates",);
+        // gen constructs get lowered to a special kind of coroutine that
+        // should directly `impl Iterator`.
+        if let ty::Coroutine(did, ..) = self_ty.kind()
+            && self.tcx().coroutine_is_gen(*did)
+        {
+            debug!(?self_ty, ?obligation, "assemble_iterator_candidates",);
 
-                candidates.vec.push(IteratorCandidate);
-            }
+            candidates.vec.push(IteratorCandidate);
+        }
+    }
+
+    fn assemble_fused_iterator_candidates(
+        &mut self,
+        obligation: &PolyTraitObligation<'tcx>,
+        candidates: &mut SelectionCandidateSet<'tcx>,
+    ) {
+        let self_ty = obligation.self_ty().skip_binder();
+        // gen constructs get lowered to a special kind of coroutine that
+        // should directly `impl FusedIterator`.
+        if let ty::Coroutine(did, ..) = self_ty.kind()
+            && self.tcx().coroutine_is_gen(*did)
+        {
+            debug!(?self_ty, ?obligation, "assemble_fused_iterator_candidates",);
+
+            candidates.vec.push(BuiltinCandidate { has_nested: false });
         }
     }
 
diff --git a/compiler/rustc_trait_selection/src/traits/select/confirmation.rs b/compiler/rustc_trait_selection/src/traits/select/confirmation.rs
index 557fc99a536..6f16194d44f 100644
--- a/compiler/rustc_trait_selection/src/traits/select/confirmation.rs
+++ b/compiler/rustc_trait_selection/src/traits/select/confirmation.rs
@@ -267,6 +267,8 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                 self.copy_clone_conditions(obligation)
             } else if Some(trait_def) == lang_items.clone_trait() {
                 self.copy_clone_conditions(obligation)
+            } else if Some(trait_def) == lang_items.fused_iterator_trait() {
+                self.fused_iterator_conditions(obligation)
             } else {
                 bug!("unexpected builtin trait {:?}", trait_def)
             };
diff --git a/compiler/rustc_trait_selection/src/traits/select/mod.rs b/compiler/rustc_trait_selection/src/traits/select/mod.rs
index 53aadfb8a44..adbc7d12a64 100644
--- a/compiler/rustc_trait_selection/src/traits/select/mod.rs
+++ b/compiler/rustc_trait_selection/src/traits/select/mod.rs
@@ -2259,6 +2259,20 @@ impl<'tcx> SelectionContext<'_, 'tcx> {
         }
     }
 
+    fn fused_iterator_conditions(
+        &mut self,
+        obligation: &PolyTraitObligation<'tcx>,
+    ) -> BuiltinImplConditions<'tcx> {
+        let self_ty = self.infcx.shallow_resolve(obligation.self_ty().skip_binder());
+        if let ty::Coroutine(did, ..) = *self_ty.kind()
+            && self.tcx().coroutine_is_gen(did)
+        {
+            BuiltinImplConditions::Where(ty::Binder::dummy(Vec::new()))
+        } else {
+            BuiltinImplConditions::None
+        }
+    }
+
     /// For default impls, we need to break apart a type into its
     /// "constituent types" -- meaning, the types that it contains.
     ///