Make sure that async closures (and fns) only capture their parent callable's parameters by move, and nothing else

2 files changed, 56 insertions, 4 deletions

diff --git a/compiler/rustc_ast_lowering/src/item.rs b/compiler/rustc_ast_lowering/src/item.rs
index 6b772c1295f..a122b4c5113 100644
--- a/compiler/rustc_ast_lowering/src/item.rs
+++ b/compiler/rustc_ast_lowering/src/item.rs
@@ -1278,10 +1278,11 @@ impl<'hir> LoweringContext<'_, 'hir> {
         };
         let closure_id = coroutine_kind.closure_id();
         let coroutine_expr = self.make_desugared_coroutine_expr(
-            // FIXME(async_closures): This should only move locals,
-            // and not upvars. Capturing closure upvars by ref doesn't
-            // work right now anyways, so whatever.
-            CaptureBy::Value { move_kw: rustc_span::DUMMY_SP },
+            // The default capture mode here is by-ref. Later on during upvar analysis,
+            // we will force the captured arguments to by-move, but for async closures,
+            // we want to make sure that we avoid unnecessarily moving captures, or else
+            // all async closures would default to `FnOnce` as their calling mode.
+            CaptureBy::Ref,
             closure_id,
             return_type_hint,
             body_span,
diff --git a/compiler/rustc_hir_typeck/src/upvar.rs b/compiler/rustc_hir_typeck/src/upvar.rs
index 82c5e566f16..864acf51025 100644
--- a/compiler/rustc_hir_typeck/src/upvar.rs
+++ b/compiler/rustc_hir_typeck/src/upvar.rs
@@ -200,6 +200,57 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             capture_information: Default::default(),
             fake_reads: Default::default(),
         };
+
+        // As noted in `lower_coroutine_body_with_moved_arguments`, we default the capture mode
+        // to `ByRef` for the `async {}` block internal to async fns/closure. This means
+        // that we would *not* be moving all of the parameters into the async block by default.
+        //
+        // We force all of these arguments to be captured by move before we do expr use analysis.
+        //
+        // FIXME(async_closures): This could be cleaned up. It's a bit janky that we're just
+        // moving all of the `LocalSource::AsyncFn` locals here.
+        if let Some(hir::CoroutineKind::Desugared(
+            _,
+            hir::CoroutineSource::Fn | hir::CoroutineSource::Closure,
+        )) = self.tcx.coroutine_kind(closure_def_id)
+        {
+            let hir::ExprKind::Block(block, _) = body.value.kind else {
+                bug!();
+            };
+            for stmt in block.stmts {
+                let hir::StmtKind::Local(hir::Local {
+                    init: Some(init),
+                    source: hir::LocalSource::AsyncFn,
+                    pat,
+                    ..
+                }) = stmt.kind
+                else {
+                    bug!();
+                };
+                let hir::PatKind::Binding(hir::BindingAnnotation(hir::ByRef::No, _), _, _, _) =
+                    pat.kind
+                else {
+                    // Complex pattern, skip the non-upvar local.
+                    continue;
+                };
+                let hir::ExprKind::Path(hir::QPath::Resolved(_, path)) = init.kind else {
+                    bug!();
+                };
+                let hir::def::Res::Local(local_id) = path.res else {
+                    bug!();
+                };
+                let place = self.place_for_root_variable(closure_def_id, local_id);
+                delegate.capture_information.push((
+                    place,
+                    ty::CaptureInfo {
+                        capture_kind_expr_id: Some(init.hir_id),
+                        path_expr_id: Some(init.hir_id),
+                        capture_kind: UpvarCapture::ByValue,
+                    },
+                ));
+            }
+        }
+
         euv::ExprUseVisitor::new(
             &mut delegate,
             &self.infcx,