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

9 files changed, 203 insertions, 60 deletions

diff --git a/compiler/rustc_hir_typeck/src/closure.rs b/compiler/rustc_hir_typeck/src/closure.rs
index 36af5394015..dbae8bfb542 100644
--- a/compiler/rustc_hir_typeck/src/closure.rs
+++ b/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,
diff --git a/compiler/rustc_hir_typeck/src/upvar.rs b/compiler/rustc_hir_typeck/src/upvar.rs
index 72e5b1ed95b..947fff67919 100644
--- a/compiler/rustc_hir_typeck/src/upvar.rs
+++ b/compiler/rustc_hir_typeck/src/upvar.rs
@@ -399,16 +399,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);
diff --git a/src/tools/miri/tests/pass/async-closure-captures.rs b/src/tools/miri/tests/pass/async-closure-captures.rs
index 3e33de32efb..cac26bfe146 100644
--- a/src/tools/miri/tests/pass/async-closure-captures.rs
+++ b/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;
+    }
 }
diff --git a/src/tools/miri/tests/pass/async-closure-captures.stderr b/src/tools/miri/tests/pass/async-closure-captures.stderr
new file mode 100644
index 00000000000..f1548aadefa
--- /dev/null
+++ b/src/tools/miri/tests/pass/async-closure-captures.stderr
@@ -0,0 +1,31 @@
+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
deleted file mode 100644
index a0db6d236fe..00000000000
--- a/src/tools/miri/tests/pass/async-closure-captures.stdout
+++ /dev/null
@@ -1,10 +0,0 @@
-Hello(0)
-Hello(0)
-Hello(1)
-Hello(1)
-Hello(2)
-Hello(3)
-Hello(3)
-Hello(4)
-Hello(4)
-Hello(5)
diff --git a/tests/ui/async-await/async-closures/captures.rs b/tests/ui/async-await/async-closures/captures.rs
index e3ab8713709..6011292b645 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
 
@@ -79,4 +79,40 @@ 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:?}");
+            //~^ 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.stderr b/tests/ui/async-await/async-closures/captures.stderr
new file mode 100644
index 00000000000..5893854e57a
--- /dev/null
+++ b/tests/ui/async-await/async-closures/captures.stderr
@@ -0,0 +1,31 @@
+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`.
diff --git a/tests/ui/async-await/async-closures/wrong-fn-kind.rs b/tests/ui/async-await/async-closures/wrong-fn-kind.rs
index 8502bb6e2d4..3d6f856874f 100644
--- a/tests/ui/async-await/async-closures/wrong-fn-kind.rs
+++ b/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() {}
diff --git a/tests/ui/async-await/async-closures/wrong-fn-kind.stderr b/tests/ui/async-await/async-closures/wrong-fn-kind.stderr
index d0f1948e48f..e56389b3202 100644
--- a/tests/ui/async-await/async-closures/wrong-fn-kind.stderr
+++ b/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`.