Auto merge of #118420 - compiler-errors:async-gen, r=eholk

Introduce support for `async gen` blocks

I'm delighted to demonstrate that `async gen` block are not very difficult to support. They're simply coroutines that yield `Poll<Option<T>>` and return `()`.

**This PR is WIP and in draft mode for now** -- I'm mostly putting it up to show folks that it's possible. This PR needs a lang-team experiment associated with it or possible an RFC, since I don't think it falls under the jurisdiction of the `gen` RFC that was recently authored by oli (https://github.com/rust-lang/rfcs/pull/3513, https://github.com/rust-lang/rust/issues/117078).

### Technical note on the pre-generator-transform yield type:

The reason that the underlying coroutines yield `Poll<Option<T>>` and not `Poll<T>` (which would make more sense, IMO, for the pre-transformed coroutine), is because the `TransformVisitor` that is used to turn coroutines into built-in state machine functions would have to destructure and reconstruct the latter into the former, which requires at least inserting a new basic block (for a `switchInt` terminator, to match on the `Poll` discriminant).

This does mean that the desugaring (at the `rustc_ast_lowering` level) of `async gen` blocks is a bit more involved. However, since we already need to intercept both `.await` and `yield` operators, I don't consider it much of a technical burden.

r? `@ghost`

4 files changed, 192 insertions, 52 deletions

diff --git a/compiler/rustc_ast_lowering/src/expr.rs b/compiler/rustc_ast_lowering/src/expr.rs
index 835ff36acbc..c287c65ff36 100644
--- a/compiler/rustc_ast_lowering/src/expr.rs
+++ b/compiler/rustc_ast_lowering/src/expr.rs
@@ -14,6 +14,7 @@ use rustc_ast::*;
 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};
@@ -196,22 +197,19 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     binder,
                     capture_clause,
                     constness,
-                    coro_kind,
+                    coroutine_kind,
                     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(
+                }) => match coroutine_kind {
+                    Some(coroutine_kind) => self.lower_expr_coroutine_closure(
                         binder,
                         *capture_clause,
                         e.id,
                         hir_id,
-                        *closure_id,
+                        *coroutine_kind,
                         fn_decl,
                         body,
                         *fn_decl_span,
@@ -325,6 +323,15 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     hir::CoroutineSource::Block,
                     |this| this.with_new_scopes(e.span, |this| this.lower_block_expr(block)),
                 ),
+                ExprKind::Gen(capture_clause, block, GenBlockKind::AsyncGen) => self
+                    .make_async_gen_expr(
+                        *capture_clause,
+                        e.id,
+                        None,
+                        e.span,
+                        hir::CoroutineSource::Block,
+                        |this| this.with_new_scopes(e.span, |this| this.lower_block_expr(block)),
+                    ),
                 ExprKind::Yield(opt_expr) => self.lower_expr_yield(e.span, opt_expr.as_deref()),
                 ExprKind::Err => hir::ExprKind::Err(
                     self.tcx.sess.span_delayed_bug(e.span, "lowered ExprKind::Err"),
@@ -736,6 +743,87 @@ impl<'hir> LoweringContext<'_, 'hir> {
         }))
     }
 
+    /// Lower a `async gen` construct to a generator that implements `AsyncIterator`.
+    ///
+    /// This results in:
+    ///
+    /// ```text
+    /// static move? |_task_context| -> () {
+    ///     <body>
+    /// }
+    /// ```
+    pub(super) fn make_async_gen_expr(
+        &mut self,
+        capture_clause: CaptureBy,
+        closure_node_id: NodeId,
+        _yield_ty: Option<hir::FnRetTy<'hir>>,
+        span: Span,
+        async_coroutine_source: hir::CoroutineSource,
+        body: impl FnOnce(&mut Self) -> hir::Expr<'hir>,
+    ) -> hir::ExprKind<'hir> {
+        let output = hir::FnRetTy::DefaultReturn(self.lower_span(span));
+
+        // Resume argument type: `ResumeTy`
+        let unstable_span = self.mark_span_with_reason(
+            DesugaringKind::Async,
+            span,
+            Some(self.allow_gen_future.clone()),
+        );
+        let resume_ty = hir::QPath::LangItem(hir::LangItem::ResumeTy, unstable_span);
+        let input_ty = hir::Ty {
+            hir_id: self.next_id(),
+            kind: hir::TyKind::Path(resume_ty),
+            span: unstable_span,
+        };
+
+        // The closure/coroutine `FnDecl` takes a single (resume) argument of type `input_ty`.
+        let fn_decl = self.arena.alloc(hir::FnDecl {
+            inputs: arena_vec![self; input_ty],
+            output,
+            c_variadic: false,
+            implicit_self: hir::ImplicitSelfKind::None,
+            lifetime_elision_allowed: false,
+        });
+
+        // Lower the argument pattern/ident. The ident is used again in the `.await` lowering.
+        let (pat, task_context_hid) = self.pat_ident_binding_mode(
+            span,
+            Ident::with_dummy_span(sym::_task_context),
+            hir::BindingAnnotation::MUT,
+        );
+        let param = hir::Param {
+            hir_id: self.next_id(),
+            pat,
+            ty_span: self.lower_span(span),
+            span: self.lower_span(span),
+        };
+        let params = arena_vec![self; param];
+
+        let body = self.lower_body(move |this| {
+            this.coroutine_kind = Some(hir::CoroutineKind::AsyncGen(async_coroutine_source));
+
+            let old_ctx = this.task_context;
+            this.task_context = Some(task_context_hid);
+            let res = body(this);
+            this.task_context = old_ctx;
+            (params, res)
+        });
+
+        // `static |_task_context| -> <ret_ty> { body }`:
+        hir::ExprKind::Closure(self.arena.alloc(hir::Closure {
+            def_id: self.local_def_id(closure_node_id),
+            binder: hir::ClosureBinder::Default,
+            capture_clause,
+            bound_generic_params: &[],
+            fn_decl,
+            body,
+            fn_decl_span: self.lower_span(span),
+            fn_arg_span: None,
+            movability: Some(hir::Movability::Static),
+            constness: hir::Constness::NotConst,
+        }))
+    }
+
     /// Forwards a possible `#[track_caller]` annotation from `outer_hir_id` to
     /// `inner_hir_id` in case the `async_fn_track_caller` feature is enabled.
     pub(super) fn maybe_forward_track_caller(
@@ -785,15 +873,18 @@ impl<'hir> LoweringContext<'_, 'hir> {
     /// ```
     fn lower_expr_await(&mut self, await_kw_span: Span, expr: &Expr) -> hir::ExprKind<'hir> {
         let full_span = expr.span.to(await_kw_span);
-        match self.coroutine_kind {
-            Some(hir::CoroutineKind::Async(_)) => {}
+
+        let is_async_gen = match self.coroutine_kind {
+            Some(hir::CoroutineKind::Async(_)) => false,
+            Some(hir::CoroutineKind::AsyncGen(_)) => true,
             Some(hir::CoroutineKind::Coroutine) | Some(hir::CoroutineKind::Gen(_)) | None => {
                 return hir::ExprKind::Err(self.tcx.sess.emit_err(AwaitOnlyInAsyncFnAndBlocks {
                     await_kw_span,
                     item_span: self.current_item,
                 }));
             }
-        }
+        };
+
         let span = self.mark_span_with_reason(DesugaringKind::Await, await_kw_span, None);
         let gen_future_span = self.mark_span_with_reason(
             DesugaringKind::Await,
@@ -882,12 +973,19 @@ impl<'hir> LoweringContext<'_, 'hir> {
             self.stmt_expr(span, match_expr)
         };
 
-        // task_context = yield ();
+        // Depending on `async` of `async gen`:
+        // async     - task_context = yield ();
+        // async gen - task_context = yield ASYNC_GEN_PENDING;
         let yield_stmt = {
-            let unit = self.expr_unit(span);
+            let yielded = if is_async_gen {
+                self.arena.alloc(self.expr_lang_item_path(span, hir::LangItem::AsyncGenPending))
+            } else {
+                self.expr_unit(span)
+            };
+
             let yield_expr = self.expr(
                 span,
-                hir::ExprKind::Yield(unit, hir::YieldSource::Await { expr: Some(expr_hir_id) }),
+                hir::ExprKind::Yield(yielded, hir::YieldSource::Await { expr: Some(expr_hir_id) }),
             );
             let yield_expr = self.arena.alloc(yield_expr);
 
@@ -997,7 +1095,11 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 }
                 Some(movability)
             }
-            Some(hir::CoroutineKind::Gen(_)) | Some(hir::CoroutineKind::Async(_)) => {
+            Some(
+                hir::CoroutineKind::Gen(_)
+                | hir::CoroutineKind::Async(_)
+                | hir::CoroutineKind::AsyncGen(_),
+            ) => {
                 panic!("non-`async`/`gen` closure body turned `async`/`gen` during lowering");
             }
             None => {
@@ -1024,18 +1126,22 @@ impl<'hir> LoweringContext<'_, 'hir> {
         (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 });
         }
@@ -1504,8 +1610,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
     }
 
     fn lower_expr_yield(&mut self, span: Span, opt_expr: Option<&Expr>) -> hir::ExprKind<'hir> {
-        match self.coroutine_kind {
-            Some(hir::CoroutineKind::Gen(_)) => {}
+        let is_async_gen = match self.coroutine_kind {
+            Some(hir::CoroutineKind::Gen(_)) => false,
+            Some(hir::CoroutineKind::AsyncGen(_)) => true,
             Some(hir::CoroutineKind::Async(_)) => {
                 return hir::ExprKind::Err(
                     self.tcx.sess.emit_err(AsyncCoroutinesNotSupported { span }),
@@ -1521,14 +1628,24 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     )
                     .emit();
                 }
-                self.coroutine_kind = Some(hir::CoroutineKind::Coroutine)
+                self.coroutine_kind = Some(hir::CoroutineKind::Coroutine);
+                false
             }
-        }
+        };
 
-        let expr =
+        let mut yielded =
             opt_expr.as_ref().map(|x| self.lower_expr(x)).unwrap_or_else(|| self.expr_unit(span));
 
-        hir::ExprKind::Yield(expr, hir::YieldSource::Yield)
+        if is_async_gen {
+            // yield async_gen_ready($expr);
+            yielded = self.expr_call_lang_item_fn(
+                span,
+                hir::LangItem::AsyncGenReady,
+                std::slice::from_ref(yielded),
+            );
+        }
+
+        hir::ExprKind::Yield(yielded, hir::YieldSource::Yield)
     }
 
     /// Desugar `ExprForLoop` from: `[opt_ident]: for <pat> in <head> <body>` into:
diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
index 80854c8a6c0..9d1f2684c39 100644
--- a/compiler/rustc_ast_lowering/src/item.rs
+++ b/compiler/rustc_ast_lowering/src/item.rs
@@ -206,19 +206,25 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     // `impl Future<Output = T>` here because lower_body
                     // only cares about the input argument patterns in the function
                     // declaration (decl), not the return types.
-                    let coro_kind = header.coro_kind;
+                    let coroutine_kind = header.coroutine_kind;
                     let body_id = this.lower_maybe_coroutine_body(
                         span,
                         hir_id,
                         decl,
-                        coro_kind,
+                        coroutine_kind,
                         body.as_deref(),
                     );
 
                     let itctx = ImplTraitContext::Universal;
                     let (generics, decl) =
                         this.lower_generics(generics, header.constness, id, &itctx, |this| {
-                            this.lower_fn_decl(decl, id, *fn_sig_span, FnDeclKind::Fn, coro_kind)
+                            this.lower_fn_decl(
+                                decl,
+                                id,
+                                *fn_sig_span,
+                                FnDeclKind::Fn,
+                                coroutine_kind,
+                            )
                         });
                     let sig = hir::FnSig {
                         decl,
@@ -734,7 +740,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     sig,
                     i.id,
                     FnDeclKind::Trait,
-                    sig.header.coro_kind,
+                    sig.header.coroutine_kind,
                 );
                 (generics, hir::TraitItemKind::Fn(sig, hir::TraitFn::Required(names)), false)
             }
@@ -743,7 +749,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     i.span,
                     hir_id,
                     &sig.decl,
-                    sig.header.coro_kind,
+                    sig.header.coroutine_kind,
                     Some(body),
                 );
                 let (generics, sig) = self.lower_method_sig(
@@ -751,7 +757,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     sig,
                     i.id,
                     FnDeclKind::Trait,
-                    sig.header.coro_kind,
+                    sig.header.coroutine_kind,
                 );
                 (generics, hir::TraitItemKind::Fn(sig, hir::TraitFn::Provided(body_id)), true)
             }
@@ -844,7 +850,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     i.span,
                     hir_id,
                     &sig.decl,
-                    sig.header.coro_kind,
+                    sig.header.coroutine_kind,
                     body.as_deref(),
                 );
                 let (generics, sig) = self.lower_method_sig(
@@ -852,7 +858,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     sig,
                     i.id,
                     if self.is_in_trait_impl { FnDeclKind::Impl } else { FnDeclKind::Inherent },
-                    sig.header.coro_kind,
+                    sig.header.coroutine_kind,
                 );
 
                 (generics, hir::ImplItemKind::Fn(sig, body_id))
@@ -1023,17 +1029,16 @@ impl<'hir> LoweringContext<'_, 'hir> {
         span: Span,
         fn_id: hir::HirId,
         decl: &FnDecl,
-        coro_kind: Option<CoroutineKind>,
+        coroutine_kind: Option<CoroutineKind>,
         body: Option<&Block>,
     ) -> hir::BodyId {
-        let (Some(coro_kind), Some(body)) = (coro_kind, body) else {
+        let (Some(coroutine_kind), Some(body)) = (coroutine_kind, body) else {
             return self.lower_fn_body_block(span, decl, body);
         };
-        let closure_id = match coro_kind {
-            CoroutineKind::Async { closure_id, .. } | CoroutineKind::Gen { closure_id, .. } => {
-                closure_id
-            }
-        };
+        // FIXME(gen_blocks): Introduce `closure_id` method and remove ALL destructuring.
+        let (CoroutineKind::Async { closure_id, .. }
+        | CoroutineKind::Gen { closure_id, .. }
+        | CoroutineKind::AsyncGen { closure_id, .. }) = coroutine_kind;
 
         self.lower_body(|this| {
             let mut parameters: Vec<hir::Param<'_>> = Vec::new();
@@ -1200,7 +1205,8 @@ impl<'hir> LoweringContext<'_, 'hir> {
 
                 this.expr_block(body)
             };
-            let coroutine_expr = match coro_kind {
+            // FIXME(gen_blocks): Consider unifying the `make_*_expr` functions.
+            let coroutine_expr = match coroutine_kind {
                 CoroutineKind::Async { .. } => this.make_async_expr(
                     CaptureBy::Value { move_kw: rustc_span::DUMMY_SP },
                     closure_id,
@@ -1217,6 +1223,14 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     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);
@@ -1233,19 +1247,19 @@ impl<'hir> LoweringContext<'_, 'hir> {
         sig: &FnSig,
         id: NodeId,
         kind: FnDeclKind,
-        coro_kind: Option<CoroutineKind>,
+        coroutine_kind: Option<CoroutineKind>,
     ) -> (&'hir hir::Generics<'hir>, hir::FnSig<'hir>) {
         let header = self.lower_fn_header(sig.header);
         let itctx = ImplTraitContext::Universal;
         let (generics, decl) =
             self.lower_generics(generics, sig.header.constness, id, &itctx, |this| {
-                this.lower_fn_decl(&sig.decl, id, sig.span, kind, coro_kind)
+                this.lower_fn_decl(&sig.decl, id, sig.span, kind, coroutine_kind)
             });
         (generics, hir::FnSig { header, decl, span: self.lower_span(sig.span) })
     }
 
     fn lower_fn_header(&mut self, h: FnHeader) -> hir::FnHeader {
-        let asyncness = if let Some(CoroutineKind::Async { span, .. }) = h.coro_kind {
+        let asyncness = if let Some(CoroutineKind::Async { span, .. }) = h.coroutine_kind {
             hir::IsAsync::Async(span)
         } else {
             hir::IsAsync::NotAsync
diff --git a/compiler/rustc_ast_lowering/src/lib.rs b/compiler/rustc_ast_lowering/src/lib.rs
index 5dda8f5a6a3..753650f7324 100644
--- a/compiler/rustc_ast_lowering/src/lib.rs
+++ b/compiler/rustc_ast_lowering/src/lib.rs
@@ -132,6 +132,7 @@ struct LoweringContext<'a, 'hir> {
 
     allow_try_trait: Lrc<[Symbol]>,
     allow_gen_future: Lrc<[Symbol]>,
+    allow_async_iterator: Lrc<[Symbol]>,
 
     /// Mapping from generics `def_id`s to TAIT generics `def_id`s.
     /// For each captured lifetime (e.g., 'a), we create a new lifetime parameter that is a generic
@@ -176,6 +177,8 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             } else {
                 [sym::gen_future].into()
             },
+            // FIXME(gen_blocks): how does `closure_track_caller`
+            allow_async_iterator: [sym::gen_future, sym::async_iterator].into(),
             generics_def_id_map: Default::default(),
             host_param_id: None,
         }
@@ -1900,13 +1903,18 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         fn_span: Span,
     ) -> hir::FnRetTy<'hir> {
         let span = self.lower_span(fn_span);
-        let opaque_ty_span = self.mark_span_with_reason(DesugaringKind::Async, span, None);
 
-        let opaque_ty_node_id = match coro {
-            CoroutineKind::Async { return_impl_trait_id, .. }
-            | CoroutineKind::Gen { return_impl_trait_id, .. } => return_impl_trait_id,
+        let (opaque_ty_node_id, allowed_features) = match coro {
+            CoroutineKind::Async { return_impl_trait_id, .. } => (return_impl_trait_id, None),
+            CoroutineKind::Gen { return_impl_trait_id, .. } => (return_impl_trait_id, None),
+            CoroutineKind::AsyncGen { return_impl_trait_id, .. } => {
+                (return_impl_trait_id, Some(self.allow_async_iterator.clone()))
+            }
         };
 
+        let opaque_ty_span =
+            self.mark_span_with_reason(DesugaringKind::Async, span, allowed_features);
+
         let captured_lifetimes: Vec<_> = self
             .resolver
             .take_extra_lifetime_params(opaque_ty_node_id)
@@ -1925,7 +1933,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
                 let bound = this.lower_coroutine_fn_output_type_to_bound(
                     output,
                     coro,
-                    span,
+                    opaque_ty_span,
                     ImplTraitContext::ReturnPositionOpaqueTy {
                         origin: hir::OpaqueTyOrigin::FnReturn(fn_def_id),
                         fn_kind,
@@ -1944,7 +1952,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
         &mut self,
         output: &FnRetTy,
         coro: CoroutineKind,
-        span: Span,
+        opaque_ty_span: Span,
         nested_impl_trait_context: ImplTraitContext,
     ) -> hir::GenericBound<'hir> {
         // Compute the `T` in `Future<Output = T>` from the return type.
@@ -1960,20 +1968,21 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
 
         // "<$assoc_ty_name = T>"
         let (assoc_ty_name, trait_lang_item) = match coro {
-            CoroutineKind::Async { .. } => (hir::FN_OUTPUT_NAME, hir::LangItem::Future),
-            CoroutineKind::Gen { .. } => (hir::ITERATOR_ITEM_NAME, hir::LangItem::Iterator),
+            CoroutineKind::Async { .. } => (sym::Output, hir::LangItem::Future),
+            CoroutineKind::Gen { .. } => (sym::Item, hir::LangItem::Iterator),
+            CoroutineKind::AsyncGen { .. } => (sym::Item, hir::LangItem::AsyncIterator),
         };
 
         let future_args = self.arena.alloc(hir::GenericArgs {
             args: &[],
-            bindings: arena_vec![self; self.assoc_ty_binding(assoc_ty_name, span, output_ty)],
+            bindings: arena_vec![self; self.assoc_ty_binding(assoc_ty_name, opaque_ty_span, output_ty)],
             parenthesized: hir::GenericArgsParentheses::No,
             span_ext: DUMMY_SP,
         });
 
         hir::GenericBound::LangItemTrait(
             trait_lang_item,
-            self.lower_span(span),
+            opaque_ty_span,
             self.next_id(),
             future_args,
         )
diff --git a/compiler/rustc_ast_lowering/src/path.rs b/compiler/rustc_ast_lowering/src/path.rs
index 7ab0805d086..efd80af5ef4 100644
--- a/compiler/rustc_ast_lowering/src/path.rs
+++ b/compiler/rustc_ast_lowering/src/path.rs
@@ -389,7 +389,7 @@ impl<'a, 'hir> LoweringContext<'a, 'hir> {
             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,