From 102b458951b8d365ec1fb206aa4653de5d61442a Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Mon, 26 Feb 2024 23:43:10 +0000
Subject: [PATCH] Support `{async closure}: Fn` in new solver

---
 .../src/solve/assembly/structural_traits.rs   | 74 ++++++++++++++++++-
 tests/ui/async-await/async-closures/is-fn.rs  | 24 ++++++
 tests/ui/async-await/async-closures/once.rs   |  6 +-
 3 files changed, 101 insertions(+), 3 deletions(-)
 create mode 100644 tests/ui/async-await/async-closures/is-fn.rs

diff --git a/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs b/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs
index 3b902dd80f59b..f65f5f9c07023 100644
--- a/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs
+++ b/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs
@@ -281,7 +281,79 @@ pub(in crate::solve) fn extract_tupled_inputs_and_output_from_callable<'tcx>(
         }
 
         // Coroutine-closures don't implement `Fn` traits the normal way.
-        ty::CoroutineClosure(..) => Err(NoSolution),
+        // Instead, they always implement `FnOnce`, but only implement
+        // `FnMut`/`Fn` if they capture no upvars, since those may borrow
+        // from the closure.
+        ty::CoroutineClosure(def_id, args) => {
+            let args = args.as_coroutine_closure();
+            let kind_ty = args.kind_ty();
+            let sig = args.coroutine_closure_sig().skip_binder();
+
+            let coroutine_ty = if let Some(closure_kind) = kind_ty.to_opt_closure_kind() {
+                if !closure_kind.extends(goal_kind) {
+                    return Err(NoSolution);
+                }
+
+                // If `Fn`/`FnMut`, we only implement this goal if we
+                // have no captures.
+                let no_borrows = match args.tupled_upvars_ty().kind() {
+                    ty::Tuple(tys) => tys.is_empty(),
+                    ty::Error(_) => false,
+                    _ => bug!("tuple_fields called on non-tuple"),
+                };
+                if closure_kind != ty::ClosureKind::FnOnce && !no_borrows {
+                    return Err(NoSolution);
+                }
+
+                sig.to_coroutine_given_kind_and_upvars(
+                    tcx,
+                    args.parent_args(),
+                    tcx.coroutine_for_closure(def_id),
+                    goal_kind,
+                    // No captures by ref, so this doesn't matter.
+                    tcx.lifetimes.re_static,
+                    args.tupled_upvars_ty(),
+                    args.coroutine_captures_by_ref_ty(),
+                )
+            } else {
+                // Closure kind is not yet determined, so we return ambiguity unless
+                // the expected kind is `FnOnce` as that is always implemented.
+                if goal_kind != ty::ClosureKind::FnOnce {
+                    return Ok(None);
+                }
+
+                let async_fn_kind_trait_def_id =
+                    tcx.require_lang_item(LangItem::AsyncFnKindHelper, None);
+                let upvars_projection_def_id = tcx
+                    .associated_items(async_fn_kind_trait_def_id)
+                    .filter_by_name_unhygienic(sym::Upvars)
+                    .next()
+                    .unwrap()
+                    .def_id;
+                let tupled_upvars_ty = Ty::new_projection(
+                    tcx,
+                    upvars_projection_def_id,
+                    [
+                        ty::GenericArg::from(kind_ty),
+                        Ty::from_closure_kind(tcx, goal_kind).into(),
+                        // No captures by ref, so this doesn't matter.
+                        tcx.lifetimes.re_static.into(),
+                        sig.tupled_inputs_ty.into(),
+                        args.tupled_upvars_ty().into(),
+                        args.coroutine_captures_by_ref_ty().into(),
+                    ],
+                );
+                sig.to_coroutine(
+                    tcx,
+                    args.parent_args(),
+                    Ty::from_closure_kind(tcx, goal_kind),
+                    tcx.coroutine_for_closure(def_id),
+                    tupled_upvars_ty,
+                )
+            };
+
+            Ok(Some(args.coroutine_closure_sig().rebind((sig.tupled_inputs_ty, coroutine_ty))))
+        }
 
         ty::Bool
         | ty::Char
diff --git a/tests/ui/async-await/async-closures/is-fn.rs b/tests/ui/async-await/async-closures/is-fn.rs
new file mode 100644
index 0000000000000..64cc28e425f52
--- /dev/null
+++ b/tests/ui/async-await/async-closures/is-fn.rs
@@ -0,0 +1,24 @@
+//@ aux-build:block-on.rs
+//@ edition:2021
+//@ build-pass
+//@ revisions: current next
+//@[next] compile-flags: -Znext-solver
+
+#![feature(async_closure)]
+
+use std::future::Future;
+
+extern crate block_on;
+
+// Check that closures that don't capture any state may implement `Fn`.
+
+fn main() {
+    block_on::block_on(async {
+        async fn call_once<F: Future>(x: impl FnOnce(&'static str) -> F) -> F::Output {
+            x("hello, world").await
+        }
+        call_once(async |x: &'static str| {
+            println!("hello, {x}");
+        }).await
+    });
+}
diff --git a/tests/ui/async-await/async-closures/once.rs b/tests/ui/async-await/async-closures/once.rs
index 55e4cedb267a9..761df3de44429 100644
--- a/tests/ui/async-await/async-closures/once.rs
+++ b/tests/ui/async-await/async-closures/once.rs
@@ -1,6 +1,8 @@
 //@ aux-build:block-on.rs
 //@ edition:2021
 //@ build-pass
+//@ revisions: current next
+//@[next] compile-flags: -Znext-solver
 
 #![feature(async_closure)]
 
@@ -8,11 +10,11 @@ use std::future::Future;
 
 extern crate block_on;
 
-struct NoCopy;
+// Check that async closures always implement `FnOnce`
 
 fn main() {
     block_on::block_on(async {
-        async fn call_once<F: Future>(x: impl Fn(&'static str) -> F) -> F::Output {
+        async fn call_once<F: Future>(x: impl FnOnce(&'static str) -> F) -> F::Output {
             x("hello, world").await
         }
         call_once(async |x: &'static str| {