Auto merge of #118457 - eholk:genfn, r=compiler-errors

Add support for `gen fn`

This builds on #116447 to add support for `gen fn` functions. For the most part we follow the same approach as desugaring `async fn`, but replacing `Future` with `Iterator` and `async {}` with `gen {}` for the body.

The version implemented here uses the return type of a `gen fn` as the yield type. For example:

```rust
gen fn count_to_three() -> i32 {
    yield 1;
    yield 2;
    yield 3;
}
```

In the future, I think we should experiment with a syntax like `gen fn count_to_three() yield i32 { ... }`, but that can go in another PR.

cc `@oli-obk` `@compiler-errors`

Author: bors

compiler/rustc_ast/src/ast.rs

@@ -1311,7 +1311,7 @@ pub struct Closure {
     pub binder: ClosureBinder,
     pub capture_clause: CaptureBy,
     pub constness: Const,
-    pub asyncness: Async,
+    pub coro_kind: Option<CoroutineKind>,
     pub movability: Movability,
     pub fn_decl: P<FnDecl>,
     pub body: P<Expr>,
@@ -2406,28 +2406,34 @@ pub enum Unsafe {
     No,
 }
 
+/// Describes what kind of coroutine markers, if any, a function has.
+///
+/// Coroutine markers are things that cause the function to generate a coroutine, such as `async`,
+/// which makes the function return `impl Future`, or `gen`, which makes the function return `impl
+/// Iterator`.
 #[derive(Copy, Clone, Encodable, Decodable, Debug)]
-pub enum Async {
-    Yes { span: Span, closure_id: NodeId, return_impl_trait_id: NodeId },
-    No,
+pub enum CoroutineKind {
+    /// `async`, which evaluates to `impl Future`
+    Async { span: Span, closure_id: NodeId, return_impl_trait_id: NodeId },
+    /// `gen`, which evaluates to `impl Iterator`
+    Gen { span: Span, closure_id: NodeId, return_impl_trait_id: NodeId },
 }
 
-#[derive(Copy, Clone, Encodable, Decodable, Debug)]
-pub enum Gen {
-    Yes { span: Span, closure_id: NodeId, return_impl_trait_id: NodeId },
-    No,
-}
-
-impl Async {
+impl CoroutineKind {
     pub fn is_async(self) -> bool {
-        matches!(self, Async::Yes { .. })
+        matches!(self, CoroutineKind::Async { .. })
+    }
+
+    pub fn is_gen(self) -> bool {
+        matches!(self, CoroutineKind::Gen { .. })
     }
 
-    /// In this case this is an `async` return, the `NodeId` for the generated `impl Trait` item.
-    pub fn opt_return_id(self) -> Option<(NodeId, Span)> {
+    /// In this case this is an `async` or `gen` return, the `NodeId` for the generated `impl Trait`
+    /// item.
+    pub fn return_id(self) -> (NodeId, Span) {
         match self {
-            Async::Yes { return_impl_trait_id, span, .. } => Some((return_impl_trait_id, span)),
-            Async::No => None,
+            CoroutineKind::Async { return_impl_trait_id, span, .. }
+            | CoroutineKind::Gen { return_impl_trait_id, span, .. } => (return_impl_trait_id, span),
         }
     }
 }
@@ -2831,8 +2837,8 @@ impl Extern {
 pub struct FnHeader {
     /// The `unsafe` keyword, if any
     pub unsafety: Unsafe,
-    /// The `async` keyword, if any
-    pub asyncness: Async,
+    /// Whether this is `async`, `gen`, or nothing.
+    pub coro_kind: Option<CoroutineKind>,
     /// The `const` keyword, if any
     pub constness: Const,
     /// The `extern` keyword and corresponding ABI string, if any
@@ -2842,22 +2848,17 @@ pub struct FnHeader {
 impl FnHeader {
     /// Does this function header have any qualifiers or is it empty?
     pub fn has_qualifiers(&self) -> bool {
-        let Self { unsafety, asyncness, constness, ext } = self;
+        let Self { unsafety, coro_kind, constness, ext } = self;
         matches!(unsafety, Unsafe::Yes(_))
-            || asyncness.is_async()
+            || coro_kind.is_some()
             || matches!(constness, Const::Yes(_))
             || !matches!(ext, Extern::None)
     }
 }
 
 impl Default for FnHeader {
     fn default() -> FnHeader {
-        FnHeader {
-            unsafety: Unsafe::No,
-            asyncness: Async::No,
-            constness: Const::No,
-            ext: Extern::None,
-        }
+        FnHeader { unsafety: Unsafe::No, coro_kind: None, constness: Const::No, ext: Extern::None }
     }
 }
 
@@ -3177,7 +3178,7 @@ mod size_asserts {
     static_assert_size!(Block, 32);
     static_assert_size!(Expr, 72);
     static_assert_size!(ExprKind, 40);
-    static_assert_size!(Fn, 152);
+    static_assert_size!(Fn, 160);
     static_assert_size!(ForeignItem, 96);
     static_assert_size!(ForeignItemKind, 24);
     static_assert_size!(GenericArg, 24);

compiler/rustc_ast/src/mut_visit.rs

@@ -121,8 +121,8 @@ pub trait MutVisitor: Sized {
         noop_visit_fn_decl(d, self);
     }
 
-    fn visit_asyncness(&mut self, a: &mut Async) {
-        noop_visit_asyncness(a, self);
+    fn visit_coro_kind(&mut self, a: &mut CoroutineKind) {
+        noop_visit_coro_kind(a, self);
     }
 
     fn visit_closure_binder(&mut self, b: &mut ClosureBinder) {
@@ -871,13 +871,14 @@ pub fn noop_visit_closure_binder<T: MutVisitor>(binder: &mut ClosureBinder, vis:
     }
 }
 
-pub fn noop_visit_asyncness<T: MutVisitor>(asyncness: &mut Async, vis: &mut T) {
-    match asyncness {
-        Async::Yes { span: _, closure_id, return_impl_trait_id } => {
+pub fn noop_visit_coro_kind<T: MutVisitor>(coro_kind: &mut CoroutineKind, vis: &mut T) {
+    match coro_kind {
+        CoroutineKind::Async { span, closure_id, return_impl_trait_id }
+        | CoroutineKind::Gen { span, closure_id, return_impl_trait_id } => {
+            vis.visit_span(span);
             vis.visit_id(closure_id);
             vis.visit_id(return_impl_trait_id);
         }
-        Async::No => {}
     }
 }
 
@@ -1170,9 +1171,9 @@ fn visit_const_item<T: MutVisitor>(
 }
 
 pub fn noop_visit_fn_header<T: MutVisitor>(header: &mut FnHeader, vis: &mut T) {
-    let FnHeader { unsafety, asyncness, constness, ext: _ } = header;
+    let FnHeader { unsafety, coro_kind, constness, ext: _ } = header;
     visit_constness(constness, vis);
-    vis.visit_asyncness(asyncness);
+    coro_kind.as_mut().map(|coro_kind| vis.visit_coro_kind(coro_kind));
     visit_unsafety(unsafety, vis);
 }
 
@@ -1406,7 +1407,7 @@ pub fn noop_visit_expr<T: MutVisitor>(
             binder,
             capture_clause,
             constness,
-            asyncness,
+            coro_kind,
             movability: _,
             fn_decl,
             body,
@@ -1415,7 +1416,7 @@ pub fn noop_visit_expr<T: MutVisitor>(
         }) => {
             vis.visit_closure_binder(binder);
             visit_constness(constness, vis);
-            vis.visit_asyncness(asyncness);
+            coro_kind.as_mut().map(|coro_kind| vis.visit_coro_kind(coro_kind));
             vis.visit_capture_by(capture_clause);
             vis.visit_fn_decl(fn_decl);
             vis.visit_expr(body);

compiler/rustc_ast/src/visit.rs

@@ -861,7 +861,7 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
         ExprKind::Closure(box Closure {
             binder,
             capture_clause,
-            asyncness: _,
+            coro_kind: _,
             constness: _,
             movability: _,
             fn_decl,

compiler/rustc_ast_lowering/src/expr.rs

@@ -195,39 +195,39 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     binder,
                     capture_clause,
                     constness,
-                    asyncness,
+                    coro_kind,
                     movability,
                     fn_decl,
                     body,
                     fn_decl_span,
                     fn_arg_span,
-                }) => {
-                    if let Async::Yes { closure_id, .. } = asyncness {
-                        self.lower_expr_async_closure(
-                            binder,
-                            *capture_clause,
-                            e.id,
-                            hir_id,
-                            *closure_id,
-                            fn_decl,
-                            body,
-                            *fn_decl_span,
-                            *fn_arg_span,
-                        )
-                    } else {
-                        self.lower_expr_closure(
-                            binder,
-                            *capture_clause,
-                            e.id,
-                            *constness,
-                            *movability,
-                            fn_decl,
-                            body,
-                            *fn_decl_span,
-                            *fn_arg_span,
-                        )
-                    }
-                }
+                }) => match coro_kind {
+                    Some(
+                        CoroutineKind::Async { closure_id, .. }
+                        | CoroutineKind::Gen { closure_id, .. },
+                    ) => self.lower_expr_async_closure(
+                        binder,
+                        *capture_clause,
+                        e.id,
+                        hir_id,
+                        *closure_id,
+                        fn_decl,
+                        body,
+                        *fn_decl_span,
+                        *fn_arg_span,
+                    ),
+                    None => self.lower_expr_closure(
+                        binder,
+                        *capture_clause,
+                        e.id,
+                        *constness,
+                        *movability,
+                        fn_decl,
+                        body,
+                        *fn_decl_span,
+                        *fn_arg_span,
+                    ),
+                },
                 ExprKind::Block(blk, opt_label) => {
                     let opt_label = self.lower_label(*opt_label);
                     hir::ExprKind::Block(self.lower_block(blk, opt_label.is_some()), opt_label)