Rollup merge of rust-lang#123350 - compiler-errors:async-closure-by-move, r=oli-obk

Actually use the inferred `ClosureKind` from signature inference in coroutine-closures

A follow-up to rust-lang#123349, which fixes another subtle bug: We were not taking into account the async closure kind we infer during closure signature inference.

When I pass a closure directly to an arg like `fn(x: impl async FnOnce())`, that should have the side-effect of artificially restricting the kind of the async closure to `ClosureKind::FnOnce`. We weren't doing this -- that's a quick fix; however, it uncovers a second, more subtle bug with the way that `move`, async closures, and `FnOnce` interact.

Specifically, when we have an async closure like:
```
let x = Struct;
let c = infer_as_fnonce(async move || {
  println!("{x:?}");
}
```

The outer closure captures `x` by move, but the inner coroutine still immutably borrows `x` from the outer closure. Since we've forced the closure to by `async FnOnce()`, we can't actually *do* a self borrow, since the signature of `AsyncFnOnce::call_once` doesn't have a borrowed lifetime. This means that all `async move` closures that are constrained to `FnOnce` will fail borrowck.

We can fix that by detecting this case specifically, and making the *inner* async closure `move` as well. This is always beneficial to closure analysis, since if we have an `async FnOnce()` that's `move`, there's no reason to ever borrow anything, so `move` isn't artificially restrictive.

Author: GuillaumeGomez

compiler/rustc_hir_typeck/src/closure.rs

@@ -227,11 +227,18 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     kind: TypeVariableOriginKind::ClosureSynthetic,
                     span: expr_span,
                 });
-                let closure_kind_ty = self.next_ty_var(TypeVariableOrigin {
-                    // FIXME(eddyb) distinguish closure kind inference variables from the rest.
-                    kind: TypeVariableOriginKind::ClosureSynthetic,
-                    span: expr_span,
-                });
+
+                let closure_kind_ty = match expected_kind {
+                    Some(kind) => Ty::from_closure_kind(tcx, kind),
+
+                    // Create a type variable (for now) to represent the closure kind.
+                    // It will be unified during the upvar inference phase (`upvar.rs`)
+                    None => self.next_ty_var(TypeVariableOrigin {
+                        kind: TypeVariableOriginKind::ClosureSynthetic,
+                        span: expr_span,
+                    }),
+                };
+
                 let coroutine_captures_by_ref_ty = self.next_ty_var(TypeVariableOrigin {
                     kind: TypeVariableOriginKind::ClosureSynthetic,
                     span: expr_span,
@@ -262,10 +269,17 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                     },
                 );
 
-                let coroutine_kind_ty = self.next_ty_var(TypeVariableOrigin {
-                    kind: TypeVariableOriginKind::ClosureSynthetic,
-                    span: expr_span,
-                });
+                let coroutine_kind_ty = match expected_kind {
+                    Some(kind) => Ty::from_coroutine_closure_kind(tcx, kind),
+
+                    // Create a type variable (for now) to represent the closure kind.
+                    // It will be unified during the upvar inference phase (`upvar.rs`)
+                    None => self.next_ty_var(TypeVariableOrigin {
+                        kind: TypeVariableOriginKind::ClosureSynthetic,
+                        span: expr_span,
+                    }),
+                };
+
                 let coroutine_upvars_ty = self.next_ty_var(TypeVariableOrigin {
                     kind: TypeVariableOriginKind::ClosureSynthetic,
                     span: expr_span,

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
@@ -399,16 +421,22 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             );
 
             // Additionally, we can now constrain the coroutine's kind type.
-            let ty::Coroutine(_, coroutine_args) =
-                *self.typeck_results.borrow().expr_ty(body.value).kind()
-            else {
-                bug!();
-            };
-            self.demand_eqtype(
-                span,
-                coroutine_args.as_coroutine().kind_ty(),
-                Ty::from_coroutine_closure_kind(self.tcx, closure_kind),
-            );
+            //
+            // We only do this if `infer_kind`, because if we have constrained
+            // the kind from closure signature inference, the kind inferred
+            // for the inner coroutine may actually be more restrictive.
+            if infer_kind {
+                let ty::Coroutine(_, coroutine_args) =
+                    *self.typeck_results.borrow().expr_ty(body.value).kind()
+                else {
+                    bug!();
+                };
+                self.demand_eqtype(
+                    span,
+                    coroutine_args.as_coroutine().kind_ty(),
+                    Ty::from_coroutine_closure_kind(self.tcx, closure_kind),
+                );
+            }
         }
 
         self.log_closure_min_capture_info(closure_def_id, span);

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));
             }
 

src/tools/miri/tests/pass/async-closure-captures.rs

@@ -88,4 +88,38 @@ async fn async_main() {
         };
         call_once(c).await;
     }
+
+    fn force_fnonce<T>(f: impl async FnOnce() -> T) -> impl async FnOnce() -> T {
+        f
+    }
+
+    // Capture something with `move`, but infer to `AsyncFnOnce`
+    {
+        let x = Hello(6);
+        let c = force_fnonce(async move || {
+            println!("{x:?}");
+        });
+        call_once(c).await;
+
+        let x = &Hello(7);
+        let c = force_fnonce(async move || {
+            println!("{x:?}");
+        });
+        call_once(c).await;
+    }
+
+    // Capture something by-ref, but infer to `AsyncFnOnce`
+    {
+        let x = Hello(8);
+        let c = force_fnonce(async || {
+            println!("{x:?}");
+        });
+        call_once(c).await;
+
+        let x = &Hello(9);
+        let c = force_fnonce(async || {
+            println!("{x:?}");
+        });
+        call_once(c).await;
+    }
 }

src/tools/miri/tests/pass/async-closure-captures.stdout

@@ -8,3 +8,7 @@ Hello(3)
 Hello(4)
 Hello(4)
 Hello(5)
+Hello(6)
+Hello(7)
+Hello(8)
+Hello(9)

tests/ui/async-await/async-closures/captures.rs

@@ -79,4 +79,38 @@ async fn async_main() {
         };
         call_once(c).await;
     }
+
+    fn force_fnonce<T>(f: impl async FnOnce() -> T) -> impl async FnOnce() -> T {
+        f
+    }
+
+    // Capture something with `move`, but infer to `AsyncFnOnce`
+    {
+        let x = Hello(6);
+        let c = force_fnonce(async move || {
+            println!("{x:?}");
+        });
+        call_once(c).await;
+
+        let x = &Hello(7);
+        let c = force_fnonce(async move || {
+            println!("{x:?}");
+        });
+        call_once(c).await;
+    }
+
+    // Capture something by-ref, but infer to `AsyncFnOnce`
+    {
+        let x = Hello(8);
+        let c = force_fnonce(async || {
+            println!("{x:?}");
+        });
+        call_once(c).await;
+
+        let x = &Hello(9);
+        let c = force_fnonce(async || {
+            println!("{x:?}");
+        });
+        call_once(c).await;
+    }
 }

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)

tests/ui/async-await/async-closures/wrong-fn-kind.rs

@@ -2,18 +2,22 @@
 
 #![feature(async_closure)]
 
-fn main() {
-    fn needs_async_fn(_: impl async Fn()) {}
+fn needs_async_fn(_: impl async Fn()) {}
 
+fn a() {
     let mut x = 1;
     needs_async_fn(async || {
-        //~^ ERROR expected a closure that implements the `async Fn` trait, but this closure only implements `async FnMut`
+        //~^ ERROR cannot borrow `x` as mutable, as it is a captured variable in a `Fn` closure
         x += 1;
     });
+}
 
+fn b() {
     let x = String::new();
     needs_async_fn(move || async move {
         //~^ ERROR expected a closure that implements the `async Fn` trait, but this closure only implements `async FnOnce`
         println!("{x}");
     });
 }
+
+fn main() {}

tests/ui/async-await/async-closures/wrong-fn-kind.stderr

@@ -1,26 +1,5 @@
-error[E0525]: expected a closure that implements the `async Fn` trait, but this closure only implements `async FnMut`
-  --> $DIR/wrong-fn-kind.rs:9:20
-   |
-LL |       needs_async_fn(async || {
-   |       -------------- -^^^^^^^
-   |       |              |
-   |  _____|______________this closure implements `async FnMut`, not `async Fn`
-   | |     |
-   | |     required by a bound introduced by this call
-LL | |
-LL | |         x += 1;
-   | |         - closure is `async FnMut` because it mutates the variable `x` here
-LL | |     });
-   | |_____- the requirement to implement `async Fn` derives from here
-   |
-note: required by a bound in `needs_async_fn`
-  --> $DIR/wrong-fn-kind.rs:6:31
-   |
-LL |     fn needs_async_fn(_: impl async Fn()) {}
-   |                               ^^^^^^^^^^ required by this bound in `needs_async_fn`
-
 error[E0525]: expected a closure that implements the `async Fn` trait, but this closure only implements `async FnOnce`
-  --> $DIR/wrong-fn-kind.rs:15:20
+  --> $DIR/wrong-fn-kind.rs:17:20
    |
 LL |       needs_async_fn(move || async move {
    |       -------------- -^^^^^^
@@ -35,11 +14,29 @@ LL | |     });
    | |_____- the requirement to implement `async Fn` derives from here
    |
 note: required by a bound in `needs_async_fn`
-  --> $DIR/wrong-fn-kind.rs:6:31
+  --> $DIR/wrong-fn-kind.rs:5:27
+   |
+LL | fn needs_async_fn(_: impl async Fn()) {}
+   |                           ^^^^^^^^^^ required by this bound in `needs_async_fn`
+
+error[E0596]: cannot borrow `x` as mutable, as it is a captured variable in a `Fn` closure
+  --> $DIR/wrong-fn-kind.rs:9:29
    |
-LL |     fn needs_async_fn(_: impl async Fn()) {}
-   |                               ^^^^^^^^^^ required by this bound in `needs_async_fn`
+LL |   fn needs_async_fn(_: impl async Fn()) {}
+   |                        --------------- change this to accept `FnMut` instead of `Fn`
+...
+LL |       needs_async_fn(async || {
+   |  _____--------------_--------_^
+   | |     |              |
+   | |     |              in this closure
+   | |     expects `Fn` instead of `FnMut`
+LL | |
+LL | |         x += 1;
+   | |         - mutable borrow occurs due to use of `x` in closure
+LL | |     });
+   | |_____^ cannot borrow as mutable
 
 error: aborting due to 2 previous errors
 
-For more information about this error, try `rustc --explain E0525`.
+Some errors have detailed explanations: E0525, E0596.
+For more information about an error, try `rustc --explain E0525`.