Add gen blocks to ast and do some broken ast lowering

Author: oli-obk

compiler/rustc_ast/src/ast.rs

@@ -1232,7 +1232,7 @@ impl Expr {
             ExprKind::Closure(..) => ExprPrecedence::Closure,
             ExprKind::Block(..) => ExprPrecedence::Block,
             ExprKind::TryBlock(..) => ExprPrecedence::TryBlock,
-            ExprKind::Async(..) => ExprPrecedence::Async,
+            ExprKind::Gen(..) => ExprPrecedence::Gen,
             ExprKind::Await(..) => ExprPrecedence::Await,
             ExprKind::Assign(..) => ExprPrecedence::Assign,
             ExprKind::AssignOp(..) => ExprPrecedence::AssignOp,
@@ -1401,11 +1401,9 @@ pub enum ExprKind {
     Closure(Box<Closure>),
     /// A block (`'label: { ... }`).
     Block(P<Block>, Option<Label>),
-    /// An async block (`async move { ... }`).
-    ///
-    /// The async block used to have a `NodeId`, which was removed in favor of
-    /// using the parent `NodeId` of the parent `Expr`.
-    Async(CaptureBy, P<Block>),
+    /// An `async` block (`async move { ... }`),
+    /// or a `gen` block (`gen move { ... }`)
+    Gen(CaptureBy, P<Block>, GenBlockKind),
     /// An await expression (`my_future.await`). Span is of await keyword.
     Await(P<Expr>, Span),
 
@@ -1469,7 +1467,6 @@ pub enum ExprKind {
 
     /// A try expression (`expr?`).
     Try(P<Expr>),
-
     /// A `yield`, with an optional value to be yielded.
     Yield(Option<P<Expr>>),
 
@@ -1495,6 +1492,28 @@ pub enum ExprKind {
     Err,
 }
 
+/// Used to differentiate between `async {}` blocks and `gen {}` blocks.
+#[derive(Clone, Encodable, Decodable, Debug, PartialEq, Eq)]
+pub enum GenBlockKind {
+    Async,
+    Gen,
+}
+
+impl fmt::Display for GenBlockKind {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.modifier().fmt(f)
+    }
+}
+
+impl GenBlockKind {
+    pub fn modifier(&self) -> &'static str {
+        match self {
+            GenBlockKind::Async => "async",
+            GenBlockKind::Gen => "gen",
+        }
+    }
+}
+
 /// The explicit `Self` type in a "qualified path". The actual
 /// path, including the trait and the associated item, is stored
 /// separately. `position` represents the index of the associated

compiler/rustc_ast/src/mut_visit.rs

@@ -1418,7 +1418,7 @@ pub fn noop_visit_expr<T: MutVisitor>(
             vis.visit_block(blk);
             visit_opt(label, |label| vis.visit_label(label));
         }
-        ExprKind::Async(_capture_by, body) => {
+        ExprKind::Gen(_capture_by, body, _) => {
             vis.visit_block(body);
         }
         ExprKind::Await(expr, await_kw_span) => {

compiler/rustc_ast/src/util/classify.rs

@@ -46,7 +46,7 @@ pub fn expr_trailing_brace(mut expr: &ast::Expr) -> Option<&ast::Expr> {
             Closure(closure) => {
                 expr = &closure.body;
             }
-            Async(..) | Block(..) | ForLoop(..) | If(..) | Loop(..) | Match(..) | Struct(..)
+            Gen(..) | Block(..) | ForLoop(..) | If(..) | Loop(..) | Match(..) | Struct(..)
             | TryBlock(..) | While(..) => break Some(expr),
             _ => break None,
         }

compiler/rustc_ast/src/util/parser.rs

@@ -285,7 +285,7 @@ pub enum ExprPrecedence {
     Block,
     TryBlock,
     Struct,
-    Async,
+    Gen,
     Await,
     Err,
 }
@@ -351,7 +351,7 @@ impl ExprPrecedence {
             | ExprPrecedence::ConstBlock
             | ExprPrecedence::Block
             | ExprPrecedence::TryBlock
-            | ExprPrecedence::Async
+            | ExprPrecedence::Gen
             | ExprPrecedence::Struct
             | ExprPrecedence::Err => PREC_PAREN,
         }

compiler/rustc_ast/src/visit.rs

@@ -872,7 +872,7 @@ pub fn walk_expr<'a, V: Visitor<'a>>(visitor: &mut V, expression: &'a Expr) {
             walk_list!(visitor, visit_label, opt_label);
             visitor.visit_block(block);
         }
-        ExprKind::Async(_, body) => {
+        ExprKind::Gen(_, body, _) => {
             visitor.visit_block(body);
         }
         ExprKind::Await(expr, _) => visitor.visit_expr(expr),

compiler/rustc_ast_lowering/src/expr.rs

@@ -183,7 +183,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     self.arena.alloc_from_iter(arms.iter().map(|x| self.lower_arm(x))),
                     hir::MatchSource::Normal,
                 ),
-                ExprKind::Async(capture_clause, block) => self.make_async_expr(
+                ExprKind::Gen(capture_clause, block, GenBlockKind::Async) => self.make_async_expr(
                     *capture_clause,
                     e.id,
                     None,
@@ -317,6 +317,14 @@ impl<'hir> LoweringContext<'_, 'hir> {
                         rest,
                     )
                 }
+                ExprKind::Gen(capture_clause, block, GenBlockKind::Gen) => self.make_gen_expr(
+                    *capture_clause,
+                    e.id,
+                    None,
+                    e.span,
+                    hir::AsyncGeneratorKind::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()),
                 ExprKind::Err => hir::ExprKind::Err(
                     self.tcx.sess.delay_span_bug(e.span, "lowered ExprKind::Err"),
@@ -661,6 +669,86 @@ impl<'hir> LoweringContext<'_, 'hir> {
         }))
     }
 
+    /// Lower a `gen` construct to a generator that implements `Future`.
+    ///
+    /// This results in:
+    ///
+    /// ```text
+    /// static move? |()| -> () {
+    ///     <body>
+    /// }
+    /// ```
+    pub(super) fn make_gen_expr(
+        &mut self,
+        capture_clause: CaptureBy,
+        closure_node_id: NodeId,
+        yield_ty: Option<hir::FnRetTy<'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,
+        body: impl FnOnce(&mut Self) -> hir::Expr<'hir>,
+    ) -> hir::ExprKind<'hir> {
+        let output = yield_ty.unwrap_or_else(|| hir::FnRetTy::DefaultReturn(self.lower_span(span)));
+
+        // Resume argument type: `ResumeTy`
+        let unstable_span =
+            self.mark_span_with_reason(DesugaringKind::Async, span, self.allow_gen_future.clone());
+        let resume_ty = hir::QPath::LangItem(hir::LangItem::ResumeTy, unstable_span, None);
+        let input_ty = hir::Ty {
+            hir_id: self.next_id(),
+            kind: hir::TyKind::Path(resume_ty),
+            span: unstable_span,
+        };
+
+        // The closure/generator `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.generator_kind = Some(hir::GeneratorKind::Iter(gen_kind));
+
+            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(

compiler/rustc_ast_pretty/src/pprust/state/expr.rs

@@ -445,8 +445,8 @@ impl<'a> State<'a> {
                 self.ibox(0);
                 self.print_block_with_attrs(blk, attrs);
             }
-            ast::ExprKind::Async(capture_clause, blk) => {
-                self.word_nbsp("async");
+            ast::ExprKind::Gen(capture_clause, blk, kind) => {
+                self.word_nbsp(kind.modifier());
                 self.print_capture_clause(*capture_clause);
                 // cbox/ibox in analogy to the `ExprKind::Block` arm above
                 self.cbox(0);

compiler/rustc_builtin_macros/src/assert/context.rs

@@ -294,7 +294,7 @@ impl<'cx, 'a> Context<'cx, 'a> {
             // sync with the `rfc-2011-nicer-assert-messages/all-expr-kinds.rs` test.
             ExprKind::Assign(_, _, _)
             | ExprKind::AssignOp(_, _, _)
-            | ExprKind::Async(_, _)
+            | ExprKind::Gen(_, _, _)
             | ExprKind::Await(_, _)
             | ExprKind::Block(_, _)
             | ExprKind::Break(_, _)

compiler/rustc_parse/src/parser/expr.rs

@@ -9,7 +9,7 @@ use super::{
 use crate::errors;
 use crate::maybe_recover_from_interpolated_ty_qpath;
 use ast::mut_visit::{noop_visit_expr, MutVisitor};
-use ast::{Path, PathSegment};
+use ast::{GenBlockKind, Path, PathSegment};
 use core::mem;
 use rustc_ast::ptr::P;
 use rustc_ast::token::{self, Delimiter, Token, TokenKind};
@@ -1422,9 +1422,6 @@ impl<'a> Parser<'a> {
             } else if this.is_try_block() {
                 this.expect_keyword(kw::Try)?;
                 this.parse_try_block(lo)
-            } else if this.is_gen_block() {
-                this.expect_keyword(kw::Gen)?;
-                this.parse_gen_block(lo)
             } else if this.eat_keyword(kw::Return) {
                 this.parse_expr_return()
             } else if this.eat_keyword(kw::Continue) {
@@ -1446,12 +1443,14 @@ impl<'a> Parser<'a> {
                 if this.check_keyword(kw::Async) {
                     if this.is_async_block() {
                         // Check for `async {` and `async move {`.
-                        this.parse_async_block()
+                        this.parse_gen_block()
                     } else {
                         this.parse_expr_closure()
                     }
                 } else if this.eat_keyword(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() { this.parse_gen_block() } else { this.parse_expr_lit() }
                 } else {
                     this.parse_expr_lit()
                 }
@@ -3043,14 +3042,6 @@ impl<'a> Parser<'a> {
         }
     }
 
-    /// Parses a `gen {...}` expression (`gen` token already eaten).
-    fn parse_gen_block(&mut self, _span_lo: Span) -> PResult<'a, P<Expr>> {
-        let (_attrs, _body) = self.parse_inner_attrs_and_block()?;
-
-        Err(errors::GenBlock { span: self.prev_token.span }
-            .into_diagnostic(&self.sess.span_diagnostic))
-    }
-
     fn is_do_catch_block(&self) -> bool {
         self.token.is_keyword(kw::Do)
             && self.is_keyword_ahead(1, &[kw::Catch])
@@ -3077,13 +3068,18 @@ impl<'a> Parser<'a> {
             && self.token.uninterpolated_span().at_least_rust_2024()
     }
 
-    /// Parses an `async move? {...}` expression.
-    fn parse_async_block(&mut self) -> PResult<'a, P<Expr>> {
+    /// Parses an `async move? {...}` or `gen move? {...}` expression.
+    fn parse_gen_block(&mut self) -> PResult<'a, P<Expr>> {
         let lo = self.token.span;
-        self.expect_keyword(kw::Async)?;
+        let kind = if self.eat_keyword(kw::Async) {
+            GenBlockKind::Async
+        } else {
+            assert!(self.eat_keyword(kw::Gen));
+            GenBlockKind::Gen
+        };
         let capture_clause = self.parse_capture_clause()?;
         let (attrs, body) = self.parse_inner_attrs_and_block()?;
-        let kind = ExprKind::Async(capture_clause, body);
+        let kind = ExprKind::Gen(capture_clause, body, kind);
         Ok(self.mk_expr_with_attrs(lo.to(self.prev_token.span), kind, attrs))
     }
 
@@ -3614,7 +3610,7 @@ impl MutVisitor for CondChecker<'_> {
             | ExprKind::Match(_, _)
             | ExprKind::Closure(_)
             | ExprKind::Block(_, _)
-            | ExprKind::Async(_, _)
+            | ExprKind::Gen(_, _, _)
             | ExprKind::TryBlock(_)
             | ExprKind::Underscore
             | ExprKind::Path(_, _)

compiler/rustc_passes/src/hir_stats.rs

@@ -567,10 +567,10 @@ impl<'v> ast_visit::Visitor<'v> for StatCollector<'v> {
             (self, e, e.kind, Id::None, ast, Expr, ExprKind),
             [
                 Array, ConstBlock, Call, MethodCall, Tup, Binary, Unary, Lit, Cast, Type, Let,
-                If, While, ForLoop, Loop, Match, Closure, Block, Async, Await, TryBlock, Assign,
+                If, While, ForLoop, Loop, Match, Closure, Block, Await, TryBlock, Assign,
                 AssignOp, Field, Index, Range, Underscore, Path, AddrOf, Break, Continue, Ret,
                 InlineAsm, FormatArgs, OffsetOf, MacCall, Struct, Repeat, Paren, Try, Yield, Yeet,
-                Become, IncludedBytes, Err
+                Become, IncludedBytes, Gen, Err
             ]
         );
         ast_visit::walk_expr(self, e)