Force the inner coroutine of an async closure to `move` if the outer closure is `move` and `FnOnce`

3 files changed, 109 insertions, 31 deletions

diff --git a/compiler/rustc_hir_typeck/src/upvar.rs b/compiler/rustc_hir_typeck/src/upvar.rs
index 9d16f0d4815..e29a410e2e5 100644
--- a/compiler/rustc_hir_typeck/src/upvar.rs
+++ b/compiler/rustc_hir_typeck/src/upvar.rs
@@ -204,6 +204,60 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             fake_reads: Default::default(),
         };
 
+        let _ = euv::ExprUseVisitor::new(
+            &FnCtxt::new(self, self.tcx.param_env(closure_def_id), closure_def_id),
+            &mut delegate,
+        )
+        .consume_body(body);
+
+        // There are several curious situations with coroutine-closures where
+        // analysis is too aggressive with borrows when the coroutine-closure is
+        // marked `move`. Specifically:
+        //
+        // 1. If 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. See the test:
+        // `async-await/async-closures/force-move-due-to-inferred-kind.rs`.
+        //
+        // 2. If the coroutine-closure is forced to be `FnOnce` due to the way it
+        // uses its upvars, but not *all* upvars would force the closure to `FnOnce`.
+        // See the test: `async-await/async-closures/force-move-due-to-actually-fnonce.rs`.
+        //
+        // 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.
+        //
+        // One *important* note is that we do a call to `process_collected_capture_information`
+        // to eagerly test whether the coroutine would end up `FnOnce`, but we do this
+        // *before* capturing all the closure args by-value below, since that would always
+        // cause the analysis to return `FnOnce`.
+        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 hir::CaptureBy::Value { move_kw } =
+                self.tcx.hir_node(parent_hir_id).expect_closure().capture_clause
+        {
+            // (1.) Closure signature inference forced this closure to `FnOnce`.
+            if let Some(ty::ClosureKind::FnOnce) = self.closure_kind(parent_ty) {
+                capture_clause = hir::CaptureBy::Value { move_kw };
+            }
+            // (2.) The way that the closure uses its upvars means it's `FnOnce`.
+            else if let (_, ty::ClosureKind::FnOnce, _) = self
+                .process_collected_capture_information(
+                    capture_clause,
+                    &delegate.capture_information,
+                )
+            {
+                capture_clause = hir::CaptureBy::Value { move_kw };
+            }
+        }
+
         // 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.
@@ -253,34 +307,6 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             }
         }
 
-        let _ = euv::ExprUseVisitor::new(
-            &FnCtxt::new(self, self.tcx.param_env(closure_def_id), closure_def_id),
-            &mut delegate,
-        )
-        .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
@@ -289,7 +315,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
         self.log_capture_analysis_first_pass(closure_def_id, &delegate.capture_information, span);
 
         let (capture_information, closure_kind, origin) = self
-            .process_collected_capture_information(capture_clause, delegate.capture_information);
+            .process_collected_capture_information(capture_clause, &delegate.capture_information);
 
         self.compute_min_captures(closure_def_id, capture_information, span);
 
@@ -545,13 +571,14 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
     fn process_collected_capture_information(
         &self,
         capture_clause: hir::CaptureBy,
-        capture_information: InferredCaptureInformation<'tcx>,
+        capture_information: &InferredCaptureInformation<'tcx>,
     ) -> (InferredCaptureInformation<'tcx>, ty::ClosureKind, Option<(Span, Place<'tcx>)>) {
         let mut closure_kind = ty::ClosureKind::LATTICE_BOTTOM;
         let mut origin: Option<(Span, Place<'tcx>)> = None;
 
         let processed = capture_information
-            .into_iter()
+            .iter()
+            .cloned()
             .map(|(place, mut capture_info)| {
                 // Apply rules for safety before inferring closure kind
                 let (place, capture_kind) =
diff --git a/tests/ui/async-await/async-closures/force-move-due-to-actually-fnonce.rs b/tests/ui/async-await/async-closures/force-move-due-to-actually-fnonce.rs
new file mode 100644
index 00000000000..ce49f55e3e3
--- /dev/null
+++ b/tests/ui/async-await/async-closures/force-move-due-to-actually-fnonce.rs
@@ -0,0 +1,27 @@
+//@ aux-build:block-on.rs
+//@ edition:2021
+//@ check-pass
+
+#![feature(async_closure)]
+
+extern crate block_on;
+
+fn consume(_: String) {}
+
+fn main() {
+    block_on::block_on(async {
+        let x = 1i32;
+        let s = String::new();
+        // `consume(s)` pulls the closure's kind down to `FnOnce`,
+        // which means that we don't treat the borrow of `x` as a
+        // self-borrow (with `'env` lifetime). This leads to a lifetime
+        // error which is solved by forcing the inner coroutine to
+        // be `move` as well, so that it moves `x`.
+        let c = async move || {
+            println!("{x}");
+            // This makes the closure FnOnce...
+            consume(s);
+        };
+        c().await;
+    });
+}
diff --git a/tests/ui/async-await/async-closures/force-move-due-to-inferred-kind.rs b/tests/ui/async-await/async-closures/force-move-due-to-inferred-kind.rs
new file mode 100644
index 00000000000..803c990ef93
--- /dev/null
+++ b/tests/ui/async-await/async-closures/force-move-due-to-inferred-kind.rs
@@ -0,0 +1,24 @@
+//@ aux-build:block-on.rs
+//@ edition:2021
+//@ check-pass
+
+#![feature(async_closure)]
+
+extern crate block_on;
+
+fn force_fnonce<T: async FnOnce()>(t: T) -> T { t }
+
+fn main() {
+    block_on::block_on(async {
+        let x = 1i32;
+        // `force_fnonce` pulls the closure's kind down to `FnOnce`,
+        // which means that we don't treat the borrow of `x` as a
+        // self-borrow (with `'env` lifetime). This leads to a lifetime
+        // error which is solved by forcing the inner coroutine to
+        // be `move` as well, so that it moves `x`.
+        let c = force_fnonce(async move || {
+            println!("{x}");
+        });
+        c().await;
+    });
+}