Support async gen fn

compiler/rustc_ast/src/ast.rs

@@ -2415,10 +2415,12 @@ pub enum Unsafe {
 /// Iterator`.
 #[derive(Copy, Clone, Encodable, Decodable, Debug)]
 pub enum CoroutineKind {
-    /// `async`, which evaluates to `impl Future`
+    /// `async`, which returns an `impl Future`
     Async { span: Span, closure_id: NodeId, return_impl_trait_id: NodeId },
-    /// `gen`, which evaluates to `impl Iterator`
+    /// `gen`, which returns an `impl Iterator`
     Gen { span: Span, closure_id: NodeId, return_impl_trait_id: NodeId },
+    /// `async gen`, which returns an `impl AsyncIterator`
+    AsyncGen { span: Span, closure_id: NodeId, return_impl_trait_id: NodeId },
 }
 
 impl CoroutineKind {
@@ -2435,7 +2437,10 @@ pub fn is_gen(self) -> bool {
     pub fn return_id(self) -> (NodeId, Span) {
         match self {
             CoroutineKind::Async { return_impl_trait_id, span, .. }
-            | CoroutineKind::Gen { return_impl_trait_id, span, .. } => (return_impl_trait_id, span),
+            | CoroutineKind::Gen { return_impl_trait_id, span, .. }
+            | CoroutineKind::AsyncGen { return_impl_trait_id, span, .. } => {
+                (return_impl_trait_id, span)
+            }
         }
     }
 }

compiler/rustc_ast/src/mut_visit.rs

@@ -874,7 +874,8 @@ pub fn noop_visit_closure_binder<T: MutVisitor>(binder: &mut ClosureBinder, vis:
 pub fn noop_visit_coroutine_kind<T: MutVisitor>(coroutine_kind: &mut CoroutineKind, vis: &mut T) {
     match coroutine_kind {
         CoroutineKind::Async { span, closure_id, return_impl_trait_id }
-        | CoroutineKind::Gen { span, closure_id, return_impl_trait_id } => {
+        | CoroutineKind::Gen { span, closure_id, return_impl_trait_id }
+        | CoroutineKind::AsyncGen { span, closure_id, return_impl_trait_id } => {
             vis.visit_span(span);
             vis.visit_id(closure_id);
             vis.visit_id(return_impl_trait_id);

compiler/rustc_ast_lowering/src/expr.rs

@@ -13,6 +13,7 @@
 use rustc_data_structures::stack::ensure_sufficient_stack;
 use rustc_hir as hir;
 use rustc_hir::def::{DefKind, Res};
+use rustc_middle::span_bug;
 use rustc_session::errors::report_lit_error;
 use rustc_span::source_map::{respan, Spanned};
 use rustc_span::symbol::{kw, sym, Ident, Symbol};
@@ -202,15 +203,12 @@ pub(super) fn lower_expr_mut(&mut self, e: &Expr) -> hir::Expr<'hir> {
                     fn_decl_span,
                     fn_arg_span,
                 }) => match coroutine_kind {
-                    Some(
+                    Some(coroutine_kind) => self.lower_expr_coroutine_closure(
-                        CoroutineKind::Async { closure_id, .. }
-                        | CoroutineKind::Gen { closure_id, .. },
-                    ) => self.lower_expr_async_closure(
                         binder,
                         *capture_clause,
                         e.id,
                         hir_id,
-                        *closure_id,
+                        *coroutine_kind,
                         fn_decl,
                         body,
                         *fn_decl_span,
@@ -1098,18 +1096,22 @@ fn lower_closure_binder<'c>(
         (binder, params)
     }
 
-    fn lower_expr_async_closure(
+    fn lower_expr_coroutine_closure(
         &mut self,
         binder: &ClosureBinder,
         capture_clause: CaptureBy,
         closure_id: NodeId,
         closure_hir_id: hir::HirId,
-        inner_closure_id: NodeId,
+        coroutine_kind: CoroutineKind,
         decl: &FnDecl,
         body: &Expr,
         fn_decl_span: Span,
         fn_arg_span: Span,
     ) -> hir::ExprKind<'hir> {
+        let CoroutineKind::Async { closure_id: inner_closure_id, .. } = coroutine_kind else {
+            span_bug!(fn_decl_span, "`async gen` and `gen` closures are not supported, yet");
+        };
+
         if let &ClosureBinder::For { span, .. } = binder {
             self.tcx.sess.emit_err(NotSupportedForLifetimeBinderAsyncClosure { span });
         }

compiler/rustc_ast_lowering/src/item.rs

@@ -1035,11 +1035,9 @@ fn lower_maybe_coroutine_body(
         let (Some(coroutine_kind), Some(body)) = (coroutine_kind, body) else {
             return self.lower_fn_body_block(span, decl, body);
         };
-        let closure_id = match coroutine_kind {
+        let (CoroutineKind::Async { closure_id, .. }
-            CoroutineKind::Async { closure_id, .. } | CoroutineKind::Gen { closure_id, .. } => {
+        | CoroutineKind::Gen { closure_id, .. }
-                closure_id
+        | CoroutineKind::AsyncGen { closure_id, .. }) = coroutine_kind;
-            }
-        };
 
         self.lower_body(|this| {
             let mut parameters: Vec<hir::Param<'_>> = Vec::new();
@@ -1224,6 +1222,14 @@ fn lower_maybe_coroutine_body(
                     hir::CoroutineSource::Fn,
                     mkbody,
                 ),
+                CoroutineKind::AsyncGen { .. } => this.make_async_gen_expr(
+                    CaptureBy::Value { move_kw: rustc_span::DUMMY_SP },
+                    closure_id,
+                    None,
+                    body.span,
+                    hir::CoroutineSource::Fn,
+                    mkbody,
+                ),
             };
 
             let hir_id = this.lower_node_id(closure_id);

compiler/rustc_ast_lowering/src/lib.rs

@@ -1904,7 +1904,8 @@ fn lower_coroutine_fn_ret_ty(
 
         let opaque_ty_node_id = match coro {
             CoroutineKind::Async { return_impl_trait_id, .. }
-            | CoroutineKind::Gen { return_impl_trait_id, .. } => return_impl_trait_id,
+            | CoroutineKind::Gen { return_impl_trait_id, .. }
+            | CoroutineKind::AsyncGen { return_impl_trait_id, .. } => return_impl_trait_id,
         };
 
         let captured_lifetimes: Vec<_> = self
@@ -1960,8 +1961,9 @@ fn lower_coroutine_fn_output_type_to_bound(
 
         // "<$assoc_ty_name = T>"
         let (assoc_ty_name, trait_lang_item) = match coro {
-            CoroutineKind::Async { .. } => (hir::FN_OUTPUT_NAME, hir::LangItem::Future),
+            CoroutineKind::Async { .. } => (sym::Output, hir::LangItem::Future),
-            CoroutineKind::Gen { .. } => (hir::ITERATOR_ITEM_NAME, hir::LangItem::Iterator),
+            CoroutineKind::Gen { .. } => (sym::Item, hir::LangItem::Iterator),
+            CoroutineKind::AsyncGen { .. } => (sym::Item, hir::LangItem::AsyncIterator),
         };
 
         let future_args = self.arena.alloc(hir::GenericArgs {

compiler/rustc_ast_lowering/src/path.rs

@@ -389,7 +389,7 @@ fn lower_parenthesized_parameter_data(
             FnRetTy::Default(_) => self.arena.alloc(self.ty_tup(*span, &[])),
         };
         let args = smallvec![GenericArg::Type(self.arena.alloc(self.ty_tup(*inputs_span, inputs)))];
-        let binding = self.assoc_ty_binding(hir::FN_OUTPUT_NAME, output_ty.span, output_ty);
+        let binding = self.assoc_ty_binding(sym::Output, output_ty.span, output_ty);
         (
             GenericArgsCtor {
                 args,

compiler/rustc_ast_pretty/src/pprust/state.rs

@@ -1498,6 +1498,10 @@ fn print_coroutine_kind(&mut self, coroutine_kind: ast::CoroutineKind) {
             ast::CoroutineKind::Async { .. } => {
                 self.word_nbsp("async");
             }
+            ast::CoroutineKind::AsyncGen { .. } => {
+                self.word_nbsp("async");
+                self.word_nbsp("gen");
+            }
         }
     }
 

compiler/rustc_hir/src/hir.rs

@@ -2255,11 +2255,6 @@ pub enum ImplItemKind<'hir> {
     Type(&'hir Ty<'hir>),
 }
 
-/// The name of the associated type for `Fn` return types.
-pub const FN_OUTPUT_NAME: Symbol = sym::Output;
-/// The name of the associated type for `Iterator` item types.
-pub const ITERATOR_ITEM_NAME: Symbol = sym::Item;
-
 /// Bind a type to an associated type (i.e., `A = Foo`).
 ///
 /// Bindings like `A: Debug` are represented as a special type `A =

compiler/rustc_parse/messages.ftl

@@ -23,8 +23,6 @@ parse_async_block_in_2015 = `async` blocks are only allowed in Rust 2018 or late
 parse_async_fn_in_2015 = `async fn` is not permitted in Rust 2015
     .label = to use `async fn`, switch to Rust 2018 or later
 
-parse_async_gen_fn = `async gen` functions are not supported
-
 parse_async_move_block_in_2015 = `async move` blocks are only allowed in Rust 2018 or later
 
 parse_async_move_order_incorrect = the order of `move` and `async` is incorrect

compiler/rustc_parse/src/errors.rs

@@ -562,13 +562,6 @@ pub(crate) struct GenFn {
     pub span: Span,
 }
 
-#[derive(Diagnostic)]
-#[diag(parse_async_gen_fn)]
-pub(crate) struct AsyncGenFn {
-    #[primary_span]
-    pub span: Span,
-}
-
 #[derive(Diagnostic)]
 #[diag(parse_comma_after_base_struct)]
 #[note]

compiler/rustc_parse/src/parser/expr.rs

@@ -1442,21 +1442,21 @@ fn parse_expr_bottom(&mut self) -> PResult<'a, P<Expr>> {
             } else if this.token.uninterpolated_span().at_least_rust_2018() {
                 // `Span:.at_least_rust_2018()` is somewhat expensive; don't get it repeatedly.
                 if this.check_keyword(kw::Async) {
-                    if this.is_gen_block(kw::Async, 0) || this.is_gen_block(kw::Gen, 1) {
+                    // FIXME(gen_blocks): Parse `gen async` and suggest swap
+                    if this.is_gen_block(kw::Async, 0) {
                         // Check for `async {` and `async move {`,
                         // or `async gen {` and `async gen move {`.
                         this.parse_gen_block()
                     } else {
                         this.parse_expr_closure()
                     }
-                } else if this.eat_keyword(kw::Await) {
+                } else if this.token.uninterpolated_span().at_least_rust_2024()
+                    && (this.is_gen_block(kw::Gen, 0)
+                        || (this.check_keyword(kw::Async) && this.is_gen_block(kw::Gen, 1)))
+                {
+                    this.parse_gen_block()
+                } else if this.eat_keyword_noexpect(kw::Await) {
                     this.recover_incorrect_await_syntax(lo, this.prev_token.span)
-                } else if this.token.uninterpolated_span().at_least_rust_2024() {
-                    if this.is_gen_block(kw::Gen, 0) {
-                        this.parse_gen_block()
-                    } else {
-                        this.parse_expr_lit()
-                    }
                 } else {
                     this.parse_expr_lit()
                 }
@@ -2235,8 +2235,8 @@ fn parse_expr_closure(&mut self) -> PResult<'a, P<Expr>> {
         let movability =
             if self.eat_keyword(kw::Static) { Movability::Static } else { Movability::Movable };
 
-        let asyncness = if self.token.uninterpolated_span().at_least_rust_2018() {
+        let coroutine_kind = if self.token.uninterpolated_span().at_least_rust_2018() {
-            self.parse_asyncness(Case::Sensitive)
+            self.parse_coroutine_kind(Case::Sensitive)
         } else {
             None
         };
@@ -2262,9 +2262,17 @@ fn parse_expr_closure(&mut self) -> PResult<'a, P<Expr>> {
             }
         };
 
-        if let Some(CoroutineKind::Async { span, .. }) = asyncness {
+        match coroutine_kind {
-            // Feature-gate `async ||` closures.
+            Some(CoroutineKind::Async { span, .. }) => {
-            self.sess.gated_spans.gate(sym::async_closure, span);
+                // Feature-gate `async ||` closures.
+                self.sess.gated_spans.gate(sym::async_closure, span);
+            }
+            Some(CoroutineKind::Gen { span, .. }) | Some(CoroutineKind::AsyncGen { span, .. }) => {
+                // Feature-gate `gen ||` and `async gen ||` closures.
+                // FIXME(gen_blocks): This perhaps should be a different gate.
+                self.sess.gated_spans.gate(sym::gen_blocks, span);
+            }
+            None => {}
         }
 
         if self.token.kind == TokenKind::Semi
@@ -2285,7 +2293,7 @@ fn parse_expr_closure(&mut self) -> PResult<'a, P<Expr>> {
                 binder,
                 capture_clause,
                 constness,
-                coroutine_kind: asyncness,
+                coroutine_kind,
                 movability,
                 fn_decl,
                 body,

compiler/rustc_parse/src/parser/item.rs

@@ -2394,10 +2394,7 @@ pub(super) fn parse_fn_front_matter(
         let constness = self.parse_constness(case);
 
         let async_start_sp = self.token.span;
-        let asyncness = self.parse_asyncness(case);
+        let coroutine_kind = self.parse_coroutine_kind(case);
-
-        let _gen_start_sp = self.token.span;
-        let genness = self.parse_genness(case);
 
         let unsafe_start_sp = self.token.span;
         let unsafety = self.parse_unsafety(case);
@@ -2405,7 +2402,7 @@ pub(super) fn parse_fn_front_matter(
         let ext_start_sp = self.token.span;
         let ext = self.parse_extern(case);
 
-        if let Some(CoroutineKind::Async { span, .. }) = asyncness {
+        if let Some(CoroutineKind::Async { span, .. }) = coroutine_kind {
             if span.is_rust_2015() {
                 self.sess.emit_err(errors::AsyncFnIn2015 {
                     span,
@@ -2414,16 +2411,11 @@ pub(super) fn parse_fn_front_matter(
             }
         }
 
-        if let Some(CoroutineKind::Gen { span, .. }) = genness {
+        match coroutine_kind {
-            self.sess.gated_spans.gate(sym::gen_blocks, span);
+            Some(CoroutineKind::Gen { span, .. }) | Some(CoroutineKind::AsyncGen { span, .. }) => {
-        }
+                self.sess.gated_spans.gate(sym::gen_blocks, span);
-
+            }
-        if let (
+            Some(CoroutineKind::Async { .. }) | None => {}
-            Some(CoroutineKind::Async { span: async_span, .. }),
-            Some(CoroutineKind::Gen { span: gen_span, .. }),
-        ) = (asyncness, genness)
-        {
-            self.sess.emit_err(errors::AsyncGenFn { span: async_span.to(gen_span) });
         }
 
         if !self.eat_keyword_case(kw::Fn, case) {
@@ -2442,7 +2434,7 @@ enum WrongKw {
 
                     // We may be able to recover
                     let mut recover_constness = constness;
-                    let mut recover_asyncness = asyncness;
+                    let mut recover_coroutine_kind = coroutine_kind;
                     let mut recover_unsafety = unsafety;
                     // This will allow the machine fix to directly place the keyword in the correct place or to indicate
                     // that the keyword is already present and the second instance should be removed.
@@ -2455,15 +2447,24 @@ enum WrongKw {
                             }
                         }
                     } else if self.check_keyword(kw::Async) {
-                        match asyncness {
+                        match coroutine_kind {
                             Some(CoroutineKind::Async { span, .. }) => {
                                 Some(WrongKw::Duplicated(span))
                             }
+                            Some(CoroutineKind::AsyncGen { span, .. }) => {
+                                Some(WrongKw::Duplicated(span))
+                            }
                             Some(CoroutineKind::Gen { .. }) => {
-                                panic!("not sure how to recover here")
+                                recover_coroutine_kind = Some(CoroutineKind::AsyncGen {
+                                    span: self.token.span,
+                                    closure_id: DUMMY_NODE_ID,
+                                    return_impl_trait_id: DUMMY_NODE_ID,
+                                });
+                                // FIXME(gen_blocks): This span is wrong, didn't want to think about it.
+                                Some(WrongKw::Misplaced(unsafe_start_sp))
                             }
                             None => {
-                                recover_asyncness = Some(CoroutineKind::Async {
+                                recover_coroutine_kind = Some(CoroutineKind::Async {
                                     span: self.token.span,
                                     closure_id: DUMMY_NODE_ID,
                                     return_impl_trait_id: DUMMY_NODE_ID,
@@ -2561,7 +2562,7 @@ enum WrongKw {
                         return Ok(FnHeader {
                             constness: recover_constness,
                             unsafety: recover_unsafety,
-                            coroutine_kind: recover_asyncness,
+                            coroutine_kind: recover_coroutine_kind,
                             ext,
                         });
                     }
@@ -2571,12 +2572,6 @@ enum WrongKw {
             }
         }
 
-        let coroutine_kind = match asyncness {
-            Some(CoroutineKind::Async { .. }) => asyncness,
-            Some(CoroutineKind::Gen { .. }) => unreachable!("asycness cannot be Gen"),
-            None => genness,
-        };
-
         Ok(FnHeader { constness, unsafety, coroutine_kind, ext })
     }
 

compiler/rustc_parse/src/parser/mod.rs

@@ -1125,23 +1125,30 @@ pub(crate) fn is_keyword_ahead(&self, dist: usize, kws: &[Symbol]) -> bool {
     }
 
     /// Parses asyncness: `async` or nothing.
-    fn parse_asyncness(&mut self, case: Case) -> Option<CoroutineKind> {
+    fn parse_coroutine_kind(&mut self, case: Case) -> Option<CoroutineKind> {
+        let span = self.token.uninterpolated_span();
         if self.eat_keyword_case(kw::Async, case) {
-            let span = self.prev_token.uninterpolated_span();
+            // FIXME(gen_blocks): Do we want to unconditionally parse `gen` and then
-            Some(CoroutineKind::Async {
+            // error if edition <= 2024, like we do with async and edition <= 2018?
-                span,
+            if self.token.uninterpolated_span().at_least_rust_2024()
-                closure_id: DUMMY_NODE_ID,
+                && self.eat_keyword_case(kw::Gen, case)
-                return_impl_trait_id: DUMMY_NODE_ID,
+            {
-            })
+                let gen_span = self.prev_token.uninterpolated_span();
-        } else {
+                Some(CoroutineKind::AsyncGen {
-            None
+                    span: span.to(gen_span),
-        }
+                    closure_id: DUMMY_NODE_ID,
-    }
+                    return_impl_trait_id: DUMMY_NODE_ID,
-
+                })
-    /// Parses genness: `gen` or nothing.
+            } else {
-    fn parse_genness(&mut self, case: Case) -> Option<CoroutineKind> {
+                Some(CoroutineKind::Async {
-        if self.token.span.at_least_rust_2024() && self.eat_keyword_case(kw::Gen, case) {
+                    span,
-            let span = self.prev_token.uninterpolated_span();
+                    closure_id: DUMMY_NODE_ID,
+                    return_impl_trait_id: DUMMY_NODE_ID,
+                })
+            }
+        } else if self.token.uninterpolated_span().at_least_rust_2024()
+            && self.eat_keyword_case(kw::Gen, case)
+        {
             Some(CoroutineKind::Gen {
                 span,
                 closure_id: DUMMY_NODE_ID,

compiler/rustc_resolve/src/def_collector.rs

@@ -287,7 +287,8 @@ fn visit_expr(&mut self, expr: &'a Expr) {
                 match closure.coroutine_kind {
                     Some(
                         CoroutineKind::Async { closure_id, .. }
-                        | CoroutineKind::Gen { closure_id, .. },
+                        | CoroutineKind::Gen { closure_id, .. }
+                        | CoroutineKind::AsyncGen { closure_id, .. },
                     ) => self.create_def(closure_id, kw::Empty, DefKind::Closure, expr.span),
                     None => closure_def,
                 }

tests/ui/coroutine/async_gen_fn.rs

@@ -1,11 +1,13 @@
 // edition: 2024
 // compile-flags: -Zunstable-options
-#![feature(gen_blocks)]
+// check-pass
 
-// async generators are not yet supported, so this test makes sure they make some kind of reasonable
+#![feature(gen_blocks, async_iterator)]
-// error.
 
-async gen fn foo() {}
+async fn bar() {}
-//~^ `async gen` functions are not supported
+
+async gen fn foo() {
+    yield bar().await;
+}
 
 fn main() {}

tests/ui/coroutine/async_gen_fn.stderr

@@ -1,8 +0,0 @@
-error: `async gen` functions are not supported
-  --> $DIR/async_gen_fn.rs:8:1
-   |
-LL | async gen fn foo() {}
-   | ^^^^^^^^^
-
-error: aborting due to 1 previous error
-