Make gen blocks implement the Iterator trait

Author: oli-obk

compiler/rustc_ast_lowering/src/expr.rs

@@ -322,7 +322,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     e.id,
                     None,
                     e.span,
-                    hir::AsyncGeneratorKind::Block,
+                    hir::AsyncCoroutineKind::Block,
                     |this| this.with_new_scopes(|this| this.lower_block_expr(block)),
                 ),
                 ExprKind::Yield(opt_expr) => self.lower_expr_yield(e.span, opt_expr.as_deref()),
@@ -686,7 +686,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         span: Span,
         // We re-use the async enum, because we have the same cases.
         // FIXME(oli-obk): rename the enum.
-        gen_kind: hir::AsyncGeneratorKind,
+        gen_kind: hir::AsyncCoroutineKind,
         body: impl FnOnce(&mut Self) -> hir::Expr<'hir>,
     ) -> hir::ExprKind<'hir> {
         let output = hir::FnRetTy::DefaultReturn(self.lower_span(span));
@@ -701,7 +701,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         });
 
         let body = self.lower_body(move |this| {
-            this.generator_kind = Some(hir::GeneratorKind::Iter(gen_kind));
+            this.coroutine_kind = Some(hir::CoroutineKind::Gen(gen_kind));
 
             let res = body(this);
             (&[], res)

compiler/rustc_hir/src/lang_items.rs

@@ -210,6 +210,7 @@ language_item_table! {
 
     FnOnceOutput,            sym::fn_once_output,      fn_once_output,             Target::AssocTy,        GenericRequirement::None;
 
+    Iterator,                sym::iterator,            iterator_trait,             Target::Trait,          GenericRequirement::Exact(0);
     Future,                  sym::future_trait,        future_trait,               Target::Trait,          GenericRequirement::Exact(0);
     CoroutineState,          sym::coroutine_state,     gen_state,                  Target::Enum,           GenericRequirement::None;
     Coroutine,               sym::coroutine,           gen_trait,                  Target::Trait,          GenericRequirement::Minimum(1);

compiler/rustc_hir_typeck/src/check.rs

@@ -59,15 +59,15 @@ pub(super) fn check_fn<'a, 'tcx>(
         && can_be_coroutine.is_some()
     {
         let yield_ty = match kind {
-            hir::GeneratorKind::Gen(..) | hir::CoroutineKind::Coroutine => {
+            hir::CoroutineKind::Gen(..) | hir::CoroutineKind::Coroutine => {
                 let yield_ty = fcx.next_ty_var(TypeVariableOrigin {
                     kind: TypeVariableOriginKind::TypeInference,
                     span,
                 });
                 fcx.require_type_is_sized(yield_ty, span, traits::SizedYieldType);
                 yield_ty
             }
-            hir::GeneratorKind::Async(..) => Ty::new_unit(tcx),
+            hir::CoroutineKind::Async(..) => Ty::new_unit(tcx),
         };
 
         // Resume type defaults to `()` if the coroutine has no argument.

compiler/rustc_hir_typeck/src/closure.rs

@@ -652,6 +652,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                         },
                     )
                 }
+                Some(hir::CoroutineKind::Gen(hir::AsyncCoroutineKind::Fn)) => todo!(),
 
                 _ => astconv.ty_infer(None, decl.output.span()),
             },

compiler/rustc_middle/src/traits/select.rs

@@ -144,6 +144,10 @@ pub enum SelectionCandidate<'tcx> {
     /// generated for an async construct.
     FutureCandidate,
 
+    /// Implementation of an `Iterator` trait by one of the generator types
+    /// generated for a gen construct.
+    IteratorCandidate,
+
     /// Implementation of a `Fn`-family trait by one of the anonymous
     /// types generated for a fn pointer type (e.g., `fn(int) -> int`)
     FnPointerCandidate {

compiler/rustc_middle/src/ty/context.rs

@@ -766,6 +766,11 @@ impl<'tcx> TyCtxt<'tcx> {
         matches!(self.coroutine_kind(def_id), Some(hir::CoroutineKind::Async(_)))
     }
 
+    /// Returns `true` if the node pointed to by `def_id` is a coroutine for a gen construct.
+    pub fn coroutine_is_gen(self, def_id: DefId) -> bool {
+        matches!(self.coroutine_kind(def_id), Some(hir::CoroutineKind::Gen(_)))
+    }
+
     pub fn stability(self) -> &'tcx stability::Index {
         self.stability_index(())
     }

compiler/rustc_span/src/symbol.rs

@@ -226,6 +226,7 @@ symbols! {
         IpAddr,
         IrTyKind,
         Is,
+        Item,
         ItemContext,
         IterEmpty,
         IterOnce,
@@ -911,6 +912,7 @@ symbols! {
         iter,
         iter_mut,
         iter_repeat,
+        iterator,
         iterator_collect_fn,
         kcfi,
         keyword,

compiler/rustc_trait_selection/src/solve/assembly/mod.rs

@@ -199,7 +199,15 @@ pub(super) trait GoalKind<'tcx>:
         goal: Goal<'tcx, Self>,
     ) -> QueryResult<'tcx>;
 
-    /// A coroutine (that doesn't come from an `async` desugaring) is known to
+    /// A coroutine (that comes from an `gen` desugaring) is known to implement
+    /// `Iterator<Item = O>`, where `O` is given by the generator's yield type
+    /// that was computed during type-checking.
+    fn consider_builtin_iterator_candidate(
+        ecx: &mut EvalCtxt<'_, 'tcx>,
+        goal: Goal<'tcx, Self>,
+    ) -> QueryResult<'tcx>;
+
+    /// A coroutine (that doesn't come from an `async` or `gen` desugaring) is known to
     /// implement `Coroutine<R, Yield = Y, Return = O>`, given the resume, yield,
     /// and return types of the coroutine computed during type-checking.
     fn consider_builtin_coroutine_candidate(
@@ -552,6 +560,8 @@ impl<'tcx> EvalCtxt<'_, 'tcx> {
             G::consider_builtin_pointee_candidate(self, goal)
         } else if lang_items.future_trait() == Some(trait_def_id) {
             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.gen_trait() == Some(trait_def_id) {
             G::consider_builtin_coroutine_candidate(self, goal)
         } else if lang_items.discriminant_kind_trait() == Some(trait_def_id) {

compiler/rustc_trait_selection/src/solve/project_goals/mod.rs

@@ -485,6 +485,37 @@ impl<'tcx> assembly::GoalKind<'tcx> for ProjectionPredicate<'tcx> {
         )
     }
 
+    fn consider_builtin_iterator_candidate(
+        ecx: &mut EvalCtxt<'_, 'tcx>,
+        goal: Goal<'tcx, Self>,
+    ) -> QueryResult<'tcx> {
+        let self_ty = goal.predicate.self_ty();
+        let ty::Coroutine(def_id, args, _) = *self_ty.kind() else {
+            return Err(NoSolution);
+        };
+
+        // Generators are not Iterators unless they come from `gen` desugaring
+        let tcx = ecx.tcx();
+        if !tcx.coroutine_is_gen(def_id) {
+            return Err(NoSolution);
+        }
+
+        let term = args.as_coroutine().yield_ty().into();
+
+        Self::consider_implied_clause(
+            ecx,
+            goal,
+            ty::ProjectionPredicate {
+                projection_ty: ty::AliasTy::new(ecx.tcx(), goal.predicate.def_id(), [self_ty]),
+                term,
+            }
+            .to_predicate(tcx),
+            // Technically, we need to check that the iterator type is Sized,
+            // but that's already proven by the generator being WF.
+            [],
+        )
+    }
+
     fn consider_builtin_coroutine_candidate(
         ecx: &mut EvalCtxt<'_, 'tcx>,
         goal: Goal<'tcx, Self>,
@@ -496,7 +527,7 @@ impl<'tcx> assembly::GoalKind<'tcx> for ProjectionPredicate<'tcx> {
 
         // `async`-desugared coroutines do not implement the coroutine trait
         let tcx = ecx.tcx();
-        if tcx.coroutine_is_async(def_id) {
+        if tcx.coroutine_is_async(def_id) || tcx.coroutine_is_gen(def_id) {
             return Err(NoSolution);
         }
 
@@ -523,7 +554,7 @@ impl<'tcx> assembly::GoalKind<'tcx> for ProjectionPredicate<'tcx> {
                 term,
             }
             .to_predicate(tcx),
-            // Technically, we need to check that the future type is Sized,
+            // Technically, we need to check that the coroutine type is Sized,
             // but that's already proven by the coroutine being WF.
             [],
         )

compiler/rustc_trait_selection/src/solve/trait_goals.rs

@@ -335,6 +335,30 @@ impl<'tcx> assembly::GoalKind<'tcx> for TraitPredicate<'tcx> {
         ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
     }
 
+    fn consider_builtin_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 `Iterator`
+        // Technically, we need to check that the iterator output type is Sized,
+        // but that's already proven by the coroutines being WF.
+        ecx.evaluate_added_goals_and_make_canonical_response(Certainty::Yes)
+    }
+
     fn consider_builtin_coroutine_candidate(
         ecx: &mut EvalCtxt<'_, 'tcx>,
         goal: Goal<'tcx, Self>,
@@ -350,7 +374,7 @@ impl<'tcx> assembly::GoalKind<'tcx> for TraitPredicate<'tcx> {
 
         // `async`-desugared coroutines do not implement the coroutine trait
         let tcx = ecx.tcx();
-        if tcx.coroutine_is_async(def_id) {
+        if tcx.coroutine_is_async(def_id) || tcx.coroutine_is_gen(def_id) {
             return Err(NoSolution);
         }