Auto merge of rust-lang#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 (rust-lang/rfcs#3513, rust-lang#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`

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 coro_kind: Option<CoroutineKind>,
+    pub coroutine_kind: Option<CoroutineKind>,
     pub movability: Movability,
     pub fn_decl: P<FnDecl>,
     pub body: P<Expr>,
@@ -1516,6 +1516,7 @@ pub enum ExprKind {
 pub enum GenBlockKind {
     Async,
     Gen,
+    AsyncGen,
 }
 
 impl fmt::Display for GenBlockKind {
@@ -1529,6 +1530,7 @@ impl GenBlockKind {
         match self {
             GenBlockKind::Async => "async",
             GenBlockKind::Gen => "gen",
+            GenBlockKind::AsyncGen => "async gen",
         }
     }
 }
@@ -2413,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 {
@@ -2433,7 +2437,10 @@ impl CoroutineKind {
     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)
+            }
         }
     }
 }
@@ -2838,7 +2845,7 @@ pub struct FnHeader {
     /// The `unsafe` keyword, if any
     pub unsafety: Unsafe,
     /// Whether this is `async`, `gen`, or nothing.
-    pub coro_kind: Option<CoroutineKind>,
+    pub coroutine_kind: Option<CoroutineKind>,
     /// The `const` keyword, if any
     pub constness: Const,
     /// The `extern` keyword and corresponding ABI string, if any
@@ -2848,17 +2855,22 @@ 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, coro_kind, constness, ext } = self;
+        let Self { unsafety, coroutine_kind, constness, ext } = self;
         matches!(unsafety, Unsafe::Yes(_))
-            || coro_kind.is_some()
+            || coroutine_kind.is_some()
             || matches!(constness, Const::Yes(_))
             || !matches!(ext, Extern::None)
     }
 }
 
 impl Default for FnHeader {
     fn default() -> FnHeader {
-        FnHeader { unsafety: Unsafe::No, coro_kind: None, constness: Const::No, ext: Extern::None }
+        FnHeader {
+            unsafety: Unsafe::No,
+            coroutine_kind: None,
+            constness: Const::No,
+            ext: Extern::None,
+        }
     }
 }
 

compiler/rustc_ast/src/mut_visit.rs

@@ -121,8 +121,8 @@ pub trait MutVisitor: Sized {
         noop_visit_fn_decl(d, self);
     }
 
-    fn visit_coro_kind(&mut self, a: &mut CoroutineKind) {
-        noop_visit_coro_kind(a, self);
+    fn visit_coroutine_kind(&mut self, a: &mut CoroutineKind) {
+        noop_visit_coroutine_kind(a, self);
     }
 
     fn visit_closure_binder(&mut self, b: &mut ClosureBinder) {
@@ -871,10 +871,11 @@ pub fn noop_visit_closure_binder<T: MutVisitor>(binder: &mut ClosureBinder, vis:
     }
 }
 
-pub fn noop_visit_coro_kind<T: MutVisitor>(coro_kind: &mut CoroutineKind, vis: &mut T) {
-    match coro_kind {
+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);
@@ -1171,9 +1172,9 @@ fn visit_const_item<T: MutVisitor>(
 }
 
 pub fn noop_visit_fn_header<T: MutVisitor>(header: &mut FnHeader, vis: &mut T) {
-    let FnHeader { unsafety, coro_kind, constness, ext: _ } = header;
+    let FnHeader { unsafety, coroutine_kind, constness, ext: _ } = header;
     visit_constness(constness, vis);
-    coro_kind.as_mut().map(|coro_kind| vis.visit_coro_kind(coro_kind));
+    coroutine_kind.as_mut().map(|coroutine_kind| vis.visit_coroutine_kind(coroutine_kind));
     visit_unsafety(unsafety, vis);
 }
 
@@ -1407,7 +1408,7 @@ pub fn noop_visit_expr<T: MutVisitor>(
             binder,
             capture_clause,
             constness,
-            coro_kind,
+            coroutine_kind,
             movability: _,
             fn_decl,
             body,
@@ -1416,7 +1417,7 @@ pub fn noop_visit_expr<T: MutVisitor>(
         }) => {
             vis.visit_closure_binder(binder);
             visit_constness(constness, vis);
-            coro_kind.as_mut().map(|coro_kind| vis.visit_coro_kind(coro_kind));
+            coroutine_kind.as_mut().map(|coroutine_kind| vis.visit_coroutine_kind(coroutine_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,
-            coro_kind: _,
+            coroutine_kind: _,
             constness: _,
             movability: _,
             fn_decl,