Auto merge of rust-lang#116447 - oli-obk:gen_fn, r=compiler-errors

Implement `gen` blocks in the 2024 edition

Coroutines tracking issue rust-lang#43122
`gen` block tracking issue rust-lang#117078

This PR implements `gen` blocks that implement `Iterator`. Most of the logic with `async` blocks is shared, and thus I renamed various types that were referring to `async` specifically.

An example usage of `gen` blocks is

```rust
fn foo() -> impl Iterator<Item = i32> {
    gen {
        yield 42;
        for i in 5..18 {
            if i.is_even() { continue }
            yield i * 2;
        }
    }
}
```

The limitations (to be resolved) of the implementation are listed in the tracking issue

Author: bors

compiler/rustc_ast/src/ast.rs

@@ -1235,7 +1235,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,
@@ -1405,11 +1405,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),
 
@@ -1499,6 +1497,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
@@ -2363,6 +2383,12 @@ pub enum Async {
     No,
 }
 
+#[derive(Copy, Clone, Encodable, Decodable, Debug)]
+pub enum Gen {
+    Yes { span: Span, closure_id: NodeId, return_impl_trait_id: NodeId },
+    No,
+}
+
 impl Async {
     pub fn is_async(self) -> bool {
         matches!(self, Async::Yes { .. })

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/token.rs

@@ -197,6 +197,7 @@ pub fn ident_can_begin_expr(name: Symbol, span: Span, is_raw: bool) -> bool {
             kw::Continue,
             kw::False,
             kw::For,
+            kw::Gen,
             kw::If,
             kw::Let,
             kw::Loop,

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::CoroutineSource::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,57 @@ impl<'hir> LoweringContext<'_, 'hir> {
         }))
     }
 
+    /// Lower a `gen` construct to a generator that implements `Iterator`.
+    ///
+    /// 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,
+        gen_kind: hir::CoroutineSource,
+        body: impl FnOnce(&mut Self) -> hir::Expr<'hir>,
+    ) -> hir::ExprKind<'hir> {
+        let output = hir::FnRetTy::DefaultReturn(self.lower_span(span));
+
+        // The closure/generator `FnDecl` takes a single (resume) argument of type `input_ty`.
+        let fn_decl = self.arena.alloc(hir::FnDecl {
+            inputs: &[],
+            output,
+            c_variadic: false,
+            implicit_self: hir::ImplicitSelfKind::None,
+            lifetime_elision_allowed: false,
+        });
+
+        let body = self.lower_body(move |this| {
+            this.coroutine_kind = Some(hir::CoroutineKind::Gen(gen_kind));
+
+            let res = body(this);
+            (&[], res)
+        });
+
+        // `static |()| -> () { 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(Movability::Movable),
+            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(
@@ -712,7 +771,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let full_span = expr.span.to(await_kw_span);
         match self.coroutine_kind {
             Some(hir::CoroutineKind::Async(_)) => {}
-            Some(hir::CoroutineKind::Coroutine) | None => {
+            Some(hir::CoroutineKind::Coroutine) | Some(hir::CoroutineKind::Gen(_)) | None => {
                 self.tcx.sess.emit_err(AwaitOnlyInAsyncFnAndBlocks {
                     await_kw_span,
                     item_span: self.current_item,
@@ -936,8 +995,8 @@ impl<'hir> LoweringContext<'_, 'hir> {
                 }
                 Some(movability)
             }
-            Some(hir::CoroutineKind::Async(_)) => {
-                panic!("non-`async` closure body turned `async` during lowering");
+            Some(hir::CoroutineKind::Gen(_)) | Some(hir::CoroutineKind::Async(_)) => {
+                panic!("non-`async`/`gen` closure body turned `async`/`gen` during lowering");
             }
             None => {
                 if movability == Movability::Static {
@@ -1445,11 +1504,22 @@ 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::Coroutine) => {}
+            Some(hir::CoroutineKind::Gen(_)) => {}
             Some(hir::CoroutineKind::Async(_)) => {
                 self.tcx.sess.emit_err(AsyncCoroutinesNotSupported { span });
             }
-            None => self.coroutine_kind = Some(hir::CoroutineKind::Coroutine),
+            Some(hir::CoroutineKind::Coroutine) | None => {
+                if !self.tcx.features().coroutines {
+                    rustc_session::parse::feature_err(
+                        &self.tcx.sess.parse_sess,
+                        sym::coroutines,
+                        span,
+                        "yield syntax is experimental",
+                    )
+                    .emit();
+                }
+                self.coroutine_kind = Some(hir::CoroutineKind::Coroutine)
+            }
         }
 
         let expr =

compiler/rustc_ast_passes/src/feature_gate.rs

@@ -554,7 +554,12 @@ pub fn check_crate(krate: &ast::Crate, sess: &Session, features: &Features) {
         "consider removing `for<...>`"
     );
     gate_all!(more_qualified_paths, "usage of qualified paths in this context is experimental");
-    gate_all!(coroutines, "yield syntax is experimental");
+    for &span in spans.get(&sym::yield_expr).iter().copied().flatten() {
+        if !span.at_least_rust_2024() {
+            gate_feature_post!(&visitor, coroutines, span, "yield syntax is experimental");
+        }
+    }
+    gate_all!(gen_blocks, "gen blocks are experimental");
     gate_all!(raw_ref_op, "raw address of syntax is experimental");
     gate_all!(const_trait_impl, "const trait impls are experimental");
     gate_all!(

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_borrowck/src/borrowck_errors.rs

@@ -373,11 +373,12 @@ impl<'cx, 'tcx> crate::MirBorrowckCtxt<'cx, 'tcx> {
         span: Span,
         yield_span: Span,
     ) -> DiagnosticBuilder<'tcx, ErrorGuaranteed> {
+        let coroutine_kind = self.body.coroutine.as_ref().unwrap().coroutine_kind;
         let mut err = struct_span_err!(
             self,
             span,
             E0626,
-            "borrow may still be in use when coroutine yields",
+            "borrow may still be in use when {coroutine_kind:#} yields",
         );
         err.span_label(yield_span, "possible yield occurs here");
         err

compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs

@@ -2491,11 +2491,17 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
 
         let (sugg_span, suggestion) = match tcx.sess.source_map().span_to_snippet(args_span) {
             Ok(string) => {
-                if string.starts_with("async ") {
-                    let pos = args_span.lo() + BytePos(6);
-                    (args_span.with_lo(pos).with_hi(pos), "move ")
-                } else if string.starts_with("async|") {
-                    let pos = args_span.lo() + BytePos(5);
+                let coro_prefix = if string.starts_with("async") {
+                    // `async` is 5 chars long. Not using `.len()` to avoid the cast from `usize` to `u32`
+                    Some(5)
+                } else if string.starts_with("gen") {
+                    // `gen` is 3 chars long
+                    Some(3)
+                } else {
+                    None
+                };
+                if let Some(n) = coro_prefix {
+                    let pos = args_span.lo() + BytePos(n);
                     (args_span.with_lo(pos).with_hi(pos), " move")
                 } else {
                     (args_span.shrink_to_lo(), "move ")
@@ -2505,6 +2511,11 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
         };
         let kind = match use_span.coroutine_kind() {
             Some(coroutine_kind) => match coroutine_kind {
+                CoroutineKind::Gen(kind) => match kind {
+                    CoroutineSource::Block => "gen block",
+                    CoroutineSource::Closure => "gen closure",
+                    _ => bug!("gen block/closure expected, but gen function found."),
+                },
                 CoroutineKind::Async(async_kind) => match async_kind {
                     CoroutineSource::Block => "async block",
                     CoroutineSource::Closure => "async closure",

compiler/rustc_borrowck/src/diagnostics/region_name.rs

@@ -698,6 +698,20 @@ impl<'tcx> MirBorrowckCtxt<'_, 'tcx> {
                             " of async function"
                         }
                     },
+                    Some(hir::CoroutineKind::Gen(gen)) => match gen {
+                        hir::CoroutineSource::Block => " of gen block",
+                        hir::CoroutineSource::Closure => " of gen closure",
+                        hir::CoroutineSource::Fn => {
+                            let parent_item =
+                                hir.get_by_def_id(hir.get_parent_item(mir_hir_id).def_id);
+                            let output = &parent_item
+                                .fn_decl()
+                                .expect("coroutine lowered from gen fn should be in fn")
+                                .output;
+                            span = output.span();
+                            " of gen function"
+                        }
+                    },
                     Some(hir::CoroutineKind::Coroutine) => " of coroutine",
                     None => " of closure",
                 };

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_codegen_ssa/src/debuginfo/type_names.rs

@@ -560,6 +560,9 @@ pub fn push_item_name(tcx: TyCtxt<'_>, def_id: DefId, qualified: bool, output: &
 
 fn coroutine_kind_label(coroutine_kind: Option<CoroutineKind>) -> &'static str {
     match coroutine_kind {
+        Some(CoroutineKind::Gen(CoroutineSource::Block)) => "gen_block",
+        Some(CoroutineKind::Gen(CoroutineSource::Closure)) => "gen_closure",
+        Some(CoroutineKind::Gen(CoroutineSource::Fn)) => "gen_fn",
         Some(CoroutineKind::Async(CoroutineSource::Block)) => "async_block",
         Some(CoroutineKind::Async(CoroutineSource::Closure)) => "async_closure",
         Some(CoroutineKind::Async(CoroutineSource::Fn)) => "async_fn",

compiler/rustc_feature/src/unstable.rs

@@ -456,6 +456,8 @@ declare_features! (
     (unstable, ffi_returns_twice, "1.34.0", Some(58314), None),
     /// Allows using `#[repr(align(...))]` on function items
     (unstable, fn_align, "1.53.0", Some(82232), None),
+    /// Allows defining gen blocks and `gen fn`.
+    (unstable, gen_blocks, "CURRENT_RUSTC_VERSION", Some(117078), None),
     /// Infer generic args for both consts and types.
     (unstable, generic_arg_infer, "1.55.0", Some(85077), None),
     /// An extension to the `generic_associated_types` feature, allowing incomplete features.