Force `move` async-closures that are `FnOnce` to make their inner coroutines also `move`

7 files changed, 57 insertions, 73 deletions

diff --git a/compiler/rustc_hir_typeck/src/upvar.rs b/compiler/rustc_hir_typeck/src/upvar.rs
index 947fff67919..c987bfb9a0e 100644
--- a/compiler/rustc_hir_typeck/src/upvar.rs
+++ b/compiler/rustc_hir_typeck/src/upvar.rs
@@ -166,7 +166,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         span: Span,
         body_id: hir::BodyId,
         body: &'tcx hir::Body<'tcx>,
-        capture_clause: hir::CaptureBy,
+        mut capture_clause: hir::CaptureBy,
     ) {
         // Extract the type of the closure.
         let ty = self.node_ty(closure_hir_id);
@@ -259,6 +259,28 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         )
         .consume_body(body);
 
+        // If a coroutine is comes from a coroutine-closure that is `move`, but
+        // the coroutine-closure was inferred to be `FnOnce` during signature
+        // inference, then it's still possible that we try to borrow upvars from
+        // the coroutine-closure because they are not used by the coroutine body
+        // in a way that forces a move.
+        //
+        // This would lead to an impossible to satisfy situation, since `AsyncFnOnce`
+        // coroutine bodies can't borrow from their parent closure. To fix this,
+        // we force the inner coroutine to also be `move`. This only matters for
+        // coroutine-closures that are `move` since otherwise they themselves will
+        // be borrowing from the outer environment, so there's no self-borrows occuring.
+        if let UpvarArgs::Coroutine(..) = args
+            && let hir::CoroutineKind::Desugared(_, hir::CoroutineSource::Closure) =
+                self.tcx.coroutine_kind(closure_def_id).expect("coroutine should have kind")
+            && let parent_hir_id =
+                self.tcx.local_def_id_to_hir_id(self.tcx.local_parent(closure_def_id))
+            && let parent_ty = self.node_ty(parent_hir_id)
+            && let Some(ty::ClosureKind::FnOnce) = self.closure_kind(parent_ty)
+        {
+            capture_clause = self.tcx.hir_node(parent_hir_id).expect_closure().capture_clause;
+        }
+
         debug!(
             "For closure={:?}, capture_information={:#?}",
             closure_def_id, delegate.capture_information
diff --git a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
index 0866205dfd0..de43f9faff9 100644
--- a/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
+++ b/compiler/rustc_mir_transform/src/coroutine/by_move_body.rs
@@ -91,15 +91,17 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
             return;
         }
 
-        let ty::Coroutine(_, coroutine_args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
-        // We don't need to generate a by-move coroutine if the kind of the coroutine is
-        // already `FnOnce` -- that means that any upvars that the closure consumes have
-        // already been taken by-value.
-        let coroutine_kind = coroutine_args.as_coroutine().kind_ty().to_opt_closure_kind().unwrap();
-        if coroutine_kind == ty::ClosureKind::FnOnce {
+        // We don't need to generate a by-move coroutine if the coroutine body was
+        // produced by the `CoroutineKindShim`, since it's already by-move.
+        if matches!(body.source.instance, ty::InstanceDef::CoroutineKindShim { .. }) {
             return;
         }
 
+        let ty::Coroutine(_, args) = *coroutine_ty.kind() else { bug!("{body:#?}") };
+        let args = args.as_coroutine();
+
+        let coroutine_kind = args.kind_ty().to_opt_closure_kind().unwrap();
+
         let parent_def_id = tcx.local_parent(coroutine_def_id);
         let ty::CoroutineClosure(_, parent_args) =
             *tcx.type_of(parent_def_id).instantiate_identity().kind()
@@ -128,6 +130,12 @@ impl<'tcx> MirPass<'tcx> for ByMoveBody {
             // the outer closure body -- we need to change the coroutine to take the
             // upvar by value.
             if coroutine_capture.is_by_ref() && !parent_capture.is_by_ref() {
+                assert_ne!(
+                    coroutine_kind,
+                    ty::ClosureKind::FnOnce,
+                    "`FnOnce` coroutine-closures return coroutines that capture from \
+                    their body; it will always result in a borrowck error!"
+                );
                 by_ref_fields.insert(FieldIdx::from_usize(num_args + idx));
             }
 
diff --git a/src/tools/miri/tests/pass/async-closure-captures.stderr b/src/tools/miri/tests/pass/async-closure-captures.stderr
deleted file mode 100644
index f1548aadefa..00000000000
--- a/src/tools/miri/tests/pass/async-closure-captures.stderr
+++ /dev/null
@@ -1,31 +0,0 @@
-error[E0597]: `x` does not live long enough
-  --> $DIR/async-closure-captures.rs:LL:CC
-   |
-LL |           let c = force_fnonce(async move || {
-   |  ____________________________________________-
-LL | |             println!("{x:?}");
-   | |                        ^ borrowed value does not live long enough
-LL | |         });
-   | |         --
-   | |         ||
-   | |         |`x` dropped here while still borrowed
-   | |_________|borrow later used here
-   |           value captured here by coroutine
-
-error[E0597]: `x` does not live long enough
-  --> $DIR/async-closure-captures.rs:LL:CC
-   |
-LL |           let c = force_fnonce(async move || {
-   |  ____________________________________________-
-LL | |             println!("{x:?}");
-   | |                        ^ borrowed value does not live long enough
-LL | |         });
-   | |         --
-   | |         ||
-   | |         |`x` dropped here while still borrowed
-   | |_________|borrow later used here
-   |           value captured here by coroutine
-
-error: aborting due to 2 previous errors
-
-For more information about this error, try `rustc --explain E0597`.
diff --git a/src/tools/miri/tests/pass/async-closure-captures.stdout b/src/tools/miri/tests/pass/async-closure-captures.stdout
new file mode 100644
index 00000000000..42a7999b2dc
--- /dev/null
+++ b/src/tools/miri/tests/pass/async-closure-captures.stdout
@@ -0,0 +1,14 @@
+Hello(0)
+Hello(0)
+Hello(1)
+Hello(1)
+Hello(2)
+Hello(3)
+Hello(3)
+Hello(4)
+Hello(4)
+Hello(5)
+Hello(6)
+Hello(7)
+Hello(8)
+Hello(9)
diff --git a/tests/ui/async-await/async-closures/captures.rs b/tests/ui/async-await/async-closures/captures.rs
index 6011292b645..0a9d0529bf5 100644
--- a/tests/ui/async-await/async-closures/captures.rs
+++ b/tests/ui/async-await/async-closures/captures.rs
@@ -1,7 +1,7 @@
 //@ aux-build:block-on.rs
 //@ edition:2021
-
-
+//@ run-pass
+//@ check-run-results
 
 // Same as miri's `tests/pass/async-closure-captures.rs`, keep in sync
 
@@ -104,14 +104,12 @@ async fn async_main() {
         let x = Hello(8);
         let c = force_fnonce(async || {
             println!("{x:?}");
-            //~^ ERROR `x` does not live long enough
         });
         call_once(c).await;
 
         let x = &Hello(9);
         let c = force_fnonce(async || {
             println!("{x:?}");
-            //~^ ERROR `x` does not live long enough
         });
         call_once(c).await;
     }
diff --git a/tests/ui/async-await/async-closures/captures.run.stdout b/tests/ui/async-await/async-closures/captures.run.stdout
index a0db6d236fe..42a7999b2dc 100644
--- a/tests/ui/async-await/async-closures/captures.run.stdout
+++ b/tests/ui/async-await/async-closures/captures.run.stdout
@@ -8,3 +8,7 @@ Hello(3)
 Hello(4)
 Hello(4)
 Hello(5)
+Hello(6)
+Hello(7)
+Hello(8)
+Hello(9)
diff --git a/tests/ui/async-await/async-closures/captures.stderr b/tests/ui/async-await/async-closures/captures.stderr
deleted file mode 100644
index 5893854e57a..00000000000
--- a/tests/ui/async-await/async-closures/captures.stderr
+++ /dev/null
@@ -1,31 +0,0 @@
-error[E0597]: `x` does not live long enough
-  --> $DIR/captures.rs:89:24
-   |
-LL |           let c = force_fnonce(async move || {
-   |  ____________________________________________-
-LL | |             println!("{x:?}");
-   | |                        ^ borrowed value does not live long enough
-LL | |         });
-   | |         --
-   | |         ||
-   | |         |`x` dropped here while still borrowed
-   | |_________|borrow later used here
-   |           value captured here by coroutine
-
-error[E0597]: `x` does not live long enough
-  --> $DIR/captures.rs:95:24
-   |
-LL |           let c = force_fnonce(async move || {
-   |  ____________________________________________-
-LL | |             println!("{x:?}");
-   | |                        ^ borrowed value does not live long enough
-LL | |         });
-   | |         --
-   | |         ||
-   | |         |`x` dropped here while still borrowed
-   | |_________|borrow later used here
-   |           value captured here by coroutine
-
-error: aborting due to 2 previous errors
-
-For more information about this error, try `rustc --explain E0597`.