Support move expressions in coroutine closures

compiler/rustc_ast_lowering/src/diagnostics.rs

@@ -148,8 +148,8 @@ pub(crate) struct ClosureCannotBeStatic {
 }
 
 #[derive(Diagnostic)]
-#[diag("`move(expr)` is only supported in plain closures")]
-pub(crate) struct MoveExprOnlyInPlainClosures {
+#[diag("`move(expr)` is only supported in closures, `async`, `gen`, and `async gen` blocks")]
+pub(crate) struct MoveExprOnlyInSupportedContexts {
     #[primary_span]
     pub span: Span,
 }

compiler/rustc_ast_lowering/src/expr.rs

@@ -21,7 +21,7 @@ mod closure;
 use crate::diagnostics::{
     AsyncCoroutinesNotSupported, AwaitOnlyInAsyncFnAndBlocks,
     FunctionalRecordUpdateDestructuringAssignment, InclusiveRangeWithNoEnd,
-    InvalidLegacyConstGenericArg, MatchArmWithNoBody, MoveExprOnlyInPlainClosures,
+    InvalidLegacyConstGenericArg, MatchArmWithNoBody, MoveExprOnlyInSupportedContexts,
     NeverPatternWithBody, NeverPatternWithGuard, UnderscoreExprLhsAssign, UseConstGenericArg,
     YieldInClosure,
 };
@@ -42,7 +42,7 @@ struct MoveExprInitializer<'a> {
     expr: &'a Expr,
 }
 
-/// State for `move(...)` expressions found while lowering one plain closure body.
+/// State for `move(...)` expressions found while lowering one closure-like body.
 pub(super) struct MoveExprState<'hir> {
     pub(super) bindings: NodeMap<(Ident, HirId)>,
     pub(super) occurrences: Vec<MoveExprOccurrence<'hir>>,
@@ -73,6 +73,12 @@ impl<'a> MoveExprInitializerFinder<'a> {
         this.visit_expr(expr);
         this.initializers
     }
+
+    fn collect_block(block: &'a Block) -> Vec<MoveExprInitializer<'a>> {
+        let mut this = Self { initializers: Vec::new() };
+        this.visit_block(block);
+        this.initializers
+    }
 }
 
 impl<'a> Visitor<'a> for MoveExprInitializerFinder<'a> {
@@ -145,13 +151,88 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let (pat, binding) = self.pat_ident(inner.span, ident);
         let Some(state) = self.move_expr_bindings.last_mut().and_then(|state| state.as_mut())
         else {
-            span_bug!(move_kw_span, "`move(...)` lowered without a plain closure body state");
+            span_bug!(move_kw_span, "`move(...)` lowered without a closure-like body state");
         };
         state.bindings.insert(id, (ident, binding));
         state.occurrences.push(MoveExprOccurrence { id, ident, pat, binding, explicit_capture });
         (ident, binding)
     }
 
+    fn lower_expr_with_move_exprs(
+        &mut self,
+        expr: hir::Expr<'hir>,
+        move_expr_state: MoveExprState<'hir>,
+        body: &Expr,
+        whole_span: Span,
+    ) -> hir::Expr<'hir> {
+        let initializers = MoveExprInitializerFinder::collect(body);
+        self.lower_expr_with_move_expr_initializers(expr, move_expr_state, initializers, whole_span)
+    }
+
+    fn lower_expr_with_move_exprs_in_block(
+        &mut self,
+        expr: hir::Expr<'hir>,
+        move_expr_state: MoveExprState<'hir>,
+        body: &Block,
+        whole_span: Span,
+    ) -> hir::Expr<'hir> {
+        let initializers = MoveExprInitializerFinder::collect_block(body);
+        self.lower_expr_with_move_expr_initializers(expr, move_expr_state, initializers, whole_span)
+    }
+
+    fn lower_expr_with_move_expr_initializers(
+        &mut self,
+        expr: hir::Expr<'hir>,
+        move_expr_state: MoveExprState<'hir>,
+        initializers: Vec<MoveExprInitializer<'_>>,
+        whole_span: Span,
+    ) -> hir::Expr<'hir> {
+        if move_expr_state.occurrences.is_empty() {
+            return expr;
+        }
+
+        let initializers = initializers
+            .into_iter()
+            .map(|initializer| (initializer.id, initializer.expr))
+            .collect::<NodeMap<_>>();
+        let mut stmts = Vec::with_capacity(move_expr_state.occurrences.len());
+        let mut initializer_bindings = NodeMap::default();
+        for occurrence in &move_expr_state.occurrences {
+            // Evaluate the expression inside `move(...)` before creating the
+            // closure/coroutine and store it in a synthetic local:
+            // `|| move(foo).bar` becomes roughly
+            // `let __move_expr_0 = foo; || __move_expr_0.bar`.
+            let expr = initializers[&occurrence.id];
+            let init = if initializer_bindings.is_empty() {
+                self.lower_expr(expr)
+            } else {
+                // Earlier entries cover nested `move(...)` expressions that
+                // appear inside this initializer, as in
+                // `move(move(foo.clone()))`.
+                let (init, _) = self.with_move_expr_bindings(
+                    Some(MoveExprState {
+                        bindings: initializer_bindings.clone(),
+                        occurrences: Vec::new(),
+                    }),
+                    |this| this.lower_expr(expr),
+                );
+                init
+            };
+            stmts.push(self.stmt_let_pat(
+                None,
+                expr.span,
+                Some(init),
+                occurrence.pat,
+                hir::LocalSource::Normal,
+            ));
+            initializer_bindings.insert(occurrence.id, (occurrence.ident, occurrence.binding));
+        }
+
+        let stmts = self.arena.alloc_from_iter(stmts);
+        let block = self.block_all(whole_span, stmts, Some(self.arena.alloc(expr)));
+        self.expr(whole_span, hir::ExprKind::Block(block, None))
+    }
+
     fn lower_exprs(&mut self, exprs: &[Box<Expr>]) -> &'hir [hir::Expr<'hir>] {
         self.arena.alloc_from_iter(exprs.iter().map(|x| self.lower_expr_mut(x)))
     }
@@ -347,7 +428,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
                     } else {
                         let guar = self
                             .dcx()
-                            .emit_err(MoveExprOnlyInPlainClosures { span: *move_kw_span });
+                            .emit_err(MoveExprOnlyInSupportedContexts { span: *move_kw_span });
                         hir::ExprKind::Err(guar)
                     }
                 }
@@ -358,23 +439,36 @@ impl<'hir> LoweringContext<'_, 'hir> {
                         GenBlockKind::Gen => hir::CoroutineDesugaring::Gen,
                         GenBlockKind::AsyncGen => hir::CoroutineDesugaring::AsyncGen,
                     };
-                    self.make_desugared_coroutine_expr(
-                        *capture_clause,
-                        e.id,
-                        None,
-                        *decl_span,
+                    let (kind, move_expr_state) =
+                        self.with_move_expr_bindings(Some(MoveExprState::default()), |this| {
+                            this.make_desugared_coroutine_expr(
+                                *capture_clause,
+                                e.id,
+                                None,
+                                *decl_span,
+                                e.span,
+                                desugaring_kind,
+                                hir::CoroutineSource::Block,
+                                |this| {
+                                    this.with_new_scopes(e.span, |this| {
+                                        this.lower_block_expr(block)
+                                    })
+                                },
+                            )
+                        });
+                    let Some(move_expr_state) = move_expr_state else {
+                        span_bug!(
+                            *decl_span,
+                            "coroutine block lowering did not return `move(...)` state"
+                        );
+                    };
+                    let expr = hir::Expr { hir_id: expr_hir_id, kind, span };
+                    return self.lower_expr_with_move_exprs_in_block(
+                        expr,
+                        move_expr_state,
+                        block,
                         e.span,
-                        desugaring_kind,
-                        hir::CoroutineSource::Block,
-                        |this| {
-                            this.with_new_scopes(e.span, |this| {
-                                let (expr, _) = this.with_move_expr_bindings(None, |this| {
-                                    this.lower_block_expr(block)
-                                });
-                                expr
-                            })
-                        },
-                    )
+                    );
                 }
                 ExprKind::Block(blk, opt_label) => {
                     // Different from loops, label of block resolves to block id rather than
@@ -886,6 +980,20 @@ impl<'hir> LoweringContext<'_, 'hir> {
             (params, res)
         });
 
+        let explicit_captures: &'hir [hir::ExplicitCapture] = if let Some(move_expr_state) =
+            self.move_expr_bindings.last().and_then(Option::as_ref)
+        {
+            self.arena.alloc_from_iter(move_expr_state.occurrences.iter().filter_map(
+                |occurrence| {
+                    occurrence
+                        .explicit_capture
+                        .then_some(hir::ExplicitCapture { var_hir_id: occurrence.binding })
+                },
+            ))
+        } else {
+            &[]
+        };
+
         // `static |<_task_context?>| -> <return_ty> { <body> }`:
         hir::ExprKind::Closure(self.arena.alloc(hir::Closure {
             def_id: closure_def_id,
@@ -898,7 +1006,7 @@ impl<'hir> LoweringContext<'_, 'hir> {
             fn_arg_span: None,
             kind: hir::ClosureKind::Coroutine(coroutine_kind),
             constness: hir::Constness::NotConst,
-            explicit_captures: &[],
+            explicit_captures,
         }))
     }
 

compiler/rustc_ast_lowering/src/expr/closure.rs

@@ -1,19 +1,18 @@
-use rustc_ast::node_id::NodeMap;
 use rustc_ast::*;
 use rustc_hir as hir;
 use rustc_hir::{HirId, Target, find_attr};
 use rustc_middle::span_bug;
 use rustc_span::Span;
 
-use super::{LoweringContext, MoveExprInitializerFinder, MoveExprState};
+use super::{LoweringContext, MoveExprState};
 use crate::FnDeclKind;
 use crate::diagnostics::{ClosureCannotBeStatic, CoroutineTooManyParameters};
 
 impl<'hir> LoweringContext<'_, 'hir> {
     // Entry point for `ExprKind::Closure`. Plain closures go through
     // `lower_expr_plain_closure_with_move_exprs`, which can wrap the lowered
-    // closure in `let` initializers for `move(...)`. Coroutine closures keep the
-    // existing coroutine-specific path and reject `move(...)` for now.
+    // closure in `let` initializers for `move(...)`. Coroutine closures use the
+    // same wrapper after building their coroutine-specific body shape.
     pub(super) fn lower_expr_closure_expr(
         &mut self,
         e: &Expr,
@@ -23,8 +22,6 @@ impl<'hir> LoweringContext<'_, 'hir> {
         let attrs = self.lower_attrs(expr_hir_id, &e.attrs, e.span, Target::from_expr(e));
 
         match closure.coroutine_kind {
-            // FIXME(TaKO8Ki): Support `move(expr)` in coroutine closures too.
-            // For the first step, we only support plain closures.
             Some(coroutine_kind) => hir::Expr {
                 hir_id: expr_hir_id,
                 kind: self.lower_expr_coroutine_closure(
@@ -117,60 +114,13 @@ impl<'hir> LoweringContext<'_, 'hir> {
             fn_arg_span,
         );
 
-        if move_expr_state.occurrences.is_empty() {
-            return hir::Expr {
-                hir_id: expr_hir_id,
-                kind: closure_kind,
-                span: self.lower_span(whole_span),
-            };
-        }
-
-        let initializers = MoveExprInitializerFinder::collect(body)
-            .into_iter()
-            .map(|initializer| (initializer.id, initializer.expr))
-            .collect::<NodeMap<_>>();
-        let mut stmts = Vec::with_capacity(move_expr_state.occurrences.len());
-        let mut initializer_bindings = NodeMap::default();
-        for occurrence in &move_expr_state.occurrences {
-            // Evaluate the expression inside `move(...)` before creating the
-            // closure and store it in a synthetic local:
-            // `|| move(foo).bar` becomes roughly
-            // `let __move_expr_0 = foo; || __move_expr_0.bar`.
-            let expr = initializers[&occurrence.id];
-            let init = if initializer_bindings.is_empty() {
-                self.lower_expr(expr)
-            } else {
-                // Earlier entries cover nested `move(...)` expressions that
-                // appear inside this initializer, as in
-                // `move(move(foo.clone()))`.
-                let (init, _) = self.with_move_expr_bindings(
-                    Some(MoveExprState {
-                        bindings: initializer_bindings.clone(),
-                        occurrences: Vec::new(),
-                    }),
-                    |this| this.lower_expr(expr),
-                );
-                init
-            };
-            stmts.push(self.stmt_let_pat(
-                None,
-                expr.span,
-                Some(init),
-                occurrence.pat,
-                hir::LocalSource::Normal,
-            ));
-            initializer_bindings.insert(occurrence.id, (occurrence.ident, occurrence.binding));
-        }
-
-        let closure_expr = self.arena.alloc(hir::Expr {
+        let closure_expr = hir::Expr {
             hir_id: expr_hir_id,
             kind: closure_kind,
             span: self.lower_span(whole_span),
-        });
+        };
 
-        let stmts = self.arena.alloc_from_iter(stmts);
-        let block = self.block_all(whole_span, stmts, Some(closure_expr));
-        self.expr(whole_span, hir::ExprKind::Block(block, None))
+        self.lower_expr_with_move_exprs(closure_expr, move_expr_state, body, whole_span)
     }
 
     // Lowers the actual plain closure node and body. The body is lowered while a
@@ -294,9 +244,9 @@ impl<'hir> LoweringContext<'_, 'hir> {
     }
 
     // Coroutine closures are lowered separately because they build a different
-    // body shape. This path pushes `None` for `move_expr_bindings`, so any
-    // `move(...)` in the coroutine body gets a targeted unsupported-position
-    // error instead of being collected like a plain closure occurrence.
+    // body shape. The source body is still lowered with `MoveExprState` active,
+    // so `move(...)` occurrences are collected and then hoisted to the outer
+    // closure body, immediately before the generated coroutine is created.
     fn lower_expr_coroutine_closure(
         &mut self,
         binder: &ClosureBinder,
@@ -329,16 +279,27 @@ impl<'hir> LoweringContext<'_, 'hir> {
             // Transform `async |x: u8| -> X { ... }` into
             // `|x: u8| || -> X { ... }`.
             let body_id = this.lower_body(|this| {
-                let ((parameters, expr), _) = this.with_move_expr_bindings(None, |this| {
-                    this.lower_coroutine_body_with_moved_arguments(
-                        &inner_decl,
-                        |this| this.with_new_scopes(fn_decl_span, |this| this.lower_expr_mut(body)),
+                let ((parameters, expr), move_expr_state) =
+                    this.with_move_expr_bindings(Some(MoveExprState::default()), |this| {
+                        this.lower_coroutine_body_with_moved_arguments(
+                            &inner_decl,
+                            |this| {
+                                this.with_new_scopes(fn_decl_span, |this| this.lower_expr_mut(body))
+                            },
+                            fn_decl_span,
+                            body.span,
+                            coroutine_kind,
+                            hir::CoroutineSource::Closure,
+                        )
+                    });
+                let Some(move_expr_state) = move_expr_state else {
+                    span_bug!(
                         fn_decl_span,
-                        body.span,
-                        coroutine_kind,
-                        hir::CoroutineSource::Closure,
-                    )
-                });
+                        "coroutine closure lowering did not return `move(...)` state"
+                    );
+                };
+
+                let expr = this.lower_expr_with_move_exprs(expr, move_expr_state, body, body.span);
 
                 this.maybe_forward_track_caller(body.span, closure_hir_id, expr.hir_id);
 

compiler/rustc_hir_analysis/src/collect.rs

@@ -1554,6 +1554,14 @@ fn coroutine_for_closure(tcx: TyCtxt<'_>, def_id: LocalDefId) -> DefId {
         bug!()
     };
 
+    let body = tcx.hir_body(body).value;
+    // `move(...)` in coroutine closures wraps the generated coroutine in an
+    // outer block of synthetic initializer lets.
+    let body = match body.kind {
+        hir::ExprKind::Block(block, None) if let Some(tail) = block.expr => tail,
+        _ => body,
+    };
+
     let &hir::Expr {
         kind:
             hir::ExprKind::Closure(&rustc_hir::Closure {
@@ -1562,7 +1570,7 @@ fn coroutine_for_closure(tcx: TyCtxt<'_>, def_id: LocalDefId) -> DefId {
                 ..
             }),
         ..
-    } = tcx.hir_body(body).value
+    } = body
     else {
         bug!()
     };