Regular closures now built-in impls for AsyncFn*

7 files changed, 318 insertions, 114 deletions

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 4fd9a29c0b2..4b95d26f9f8 100644
--- a/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs
+++ b/compiler/rustc_trait_selection/src/solve/assembly/structural_traits.rs
@@ -394,7 +394,78 @@ pub(in crate::solve) fn extract_tupled_inputs_and_output_from_async_callable<'tc
             ))
         }
 
-        ty::FnDef(..) | ty::FnPtr(..) | ty::Closure(..) => Err(NoSolution),
+        ty::FnDef(..) | ty::FnPtr(..) => {
+            let bound_sig = self_ty.fn_sig(tcx);
+            let sig = bound_sig.skip_binder();
+            let future_trait_def_id = tcx.require_lang_item(LangItem::Future, None);
+            // `FnDef` and `FnPtr` only implement `AsyncFn*` when their
+            // return type implements `Future`.
+            let nested = vec![
+                bound_sig
+                    .rebind(ty::TraitRef::new(tcx, future_trait_def_id, [sig.output()]))
+                    .to_predicate(tcx),
+            ];
+            let future_output_def_id = tcx
+                .associated_items(future_trait_def_id)
+                .filter_by_name_unhygienic(sym::Output)
+                .next()
+                .unwrap()
+                .def_id;
+            let future_output_ty = Ty::new_projection(tcx, future_output_def_id, [sig.output()]);
+            Ok((
+                bound_sig.rebind((Ty::new_tup(tcx, sig.inputs()), sig.output(), future_output_ty)),
+                nested,
+            ))
+        }
+        ty::Closure(_, args) => {
+            let args = args.as_closure();
+            let bound_sig = args.sig();
+            let sig = bound_sig.skip_binder();
+            let future_trait_def_id = tcx.require_lang_item(LangItem::Future, None);
+            // `Closure`s only implement `AsyncFn*` when their return type
+            // implements `Future`.
+            let mut nested = vec![
+                bound_sig
+                    .rebind(ty::TraitRef::new(tcx, future_trait_def_id, [sig.output()]))
+                    .to_predicate(tcx),
+            ];
+
+            // Additionally, we need to check that the closure kind
+            // is still compatible.
+            let kind_ty = args.kind_ty();
+            if let Some(closure_kind) = kind_ty.to_opt_closure_kind() {
+                if !closure_kind.extends(goal_kind) {
+                    return Err(NoSolution);
+                }
+            } else {
+                let async_fn_kind_trait_def_id =
+                    tcx.require_lang_item(LangItem::AsyncFnKindHelper, None);
+                // When we don't know the closure kind (and therefore also the closure's upvars,
+                // which are computed at the same time), we must delay the computation of the
+                // generator's upvars. We do this using the `AsyncFnKindHelper`, which as a trait
+                // goal functions similarly to the old `ClosureKind` predicate, and ensures that
+                // the goal kind <= the closure kind. As a projection `AsyncFnKindHelper::Upvars`
+                // will project to the right upvars for the generator, appending the inputs and
+                // coroutine upvars respecting the closure kind.
+                nested.push(
+                    ty::TraitRef::new(
+                        tcx,
+                        async_fn_kind_trait_def_id,
+                        [kind_ty, Ty::from_closure_kind(tcx, goal_kind)],
+                    )
+                    .to_predicate(tcx),
+                );
+            }
+
+            let future_output_def_id = tcx
+                .associated_items(future_trait_def_id)
+                .filter_by_name_unhygienic(sym::Output)
+                .next()
+                .unwrap()
+                .def_id;
+            let future_output_ty = Ty::new_projection(tcx, future_output_def_id, [sig.output()]);
+            Ok((bound_sig.rebind((sig.inputs()[0], sig.output(), future_output_ty)), nested))
+        }
 
         ty::Bool
         | ty::Char
diff --git a/compiler/rustc_trait_selection/src/traits/project.rs b/compiler/rustc_trait_selection/src/traits/project.rs
index 88c28761d25..f45a20ccd32 100644
--- a/compiler/rustc_trait_selection/src/traits/project.rs
+++ b/compiler/rustc_trait_selection/src/traits/project.rs
@@ -2450,14 +2450,6 @@ fn confirm_async_closure_candidate<'cx, 'tcx>(
 ) -> Progress<'tcx> {
     let tcx = selcx.tcx();
     let self_ty = selcx.infcx.shallow_resolve(obligation.predicate.self_ty());
-    let ty::CoroutineClosure(def_id, args) = *self_ty.kind() else {
-        unreachable!(
-            "expected coroutine-closure self type for coroutine-closure candidate, found {self_ty}"
-        )
-    };
-    let args = args.as_coroutine_closure();
-    let kind_ty = args.kind_ty();
-    let sig = args.coroutine_closure_sig().skip_binder();
 
     let goal_kind =
         tcx.async_fn_trait_kind_from_def_id(obligation.predicate.trait_def_id(tcx)).unwrap();
@@ -2465,84 +2457,163 @@ fn confirm_async_closure_candidate<'cx, 'tcx>(
         ty::ClosureKind::Fn | ty::ClosureKind::FnMut => obligation.predicate.args.region_at(2),
         ty::ClosureKind::FnOnce => tcx.lifetimes.re_static,
     };
-
     let item_name = tcx.item_name(obligation.predicate.def_id);
-    let term = match item_name {
-        sym::CallOnceFuture | sym::CallMutFuture | sym::CallFuture => {
-            if let Some(closure_kind) = kind_ty.to_opt_closure_kind() {
-                if !closure_kind.extends(goal_kind) {
-                    bug!("we should not be confirming if the closure kind is not met");
+
+    let poly_cache_entry = match *self_ty.kind() {
+        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 term = match item_name {
+                sym::CallOnceFuture | sym::CallMutFuture | sym::CallFuture => {
+                    if let Some(closure_kind) = kind_ty.to_opt_closure_kind() {
+                        if !closure_kind.extends(goal_kind) {
+                            bug!("we should not be confirming if the closure kind is not met");
+                        }
+                        sig.to_coroutine_given_kind_and_upvars(
+                            tcx,
+                            args.parent_args(),
+                            tcx.coroutine_for_closure(def_id),
+                            goal_kind,
+                            env_region,
+                            args.tupled_upvars_ty(),
+                            args.coroutine_captures_by_ref_ty(),
+                        )
+                    } else {
+                        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;
+                        // When we don't know the closure kind (and therefore also the closure's upvars,
+                        // which are computed at the same time), we must delay the computation of the
+                        // generator's upvars. We do this using the `AsyncFnKindHelper`, which as a trait
+                        // goal functions similarly to the old `ClosureKind` predicate, and ensures that
+                        // the goal kind <= the closure kind. As a projection `AsyncFnKindHelper::Upvars`
+                        // will project to the right upvars for the generator, appending the inputs and
+                        // coroutine upvars respecting the closure kind.
+                        // N.B. No need to register a `AsyncFnKindHelper` goal here, it's already in `nested`.
+                        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(),
+                                env_region.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,
+                        )
+                    }
                 }
-                sig.to_coroutine_given_kind_and_upvars(
+                sym::Output => sig.return_ty,
+                name => bug!("no such associated type: {name}"),
+            };
+            let projection_ty = match item_name {
+                sym::CallOnceFuture | sym::Output => ty::AliasTy::new(
                     tcx,
-                    args.parent_args(),
-                    tcx.coroutine_for_closure(def_id),
-                    goal_kind,
-                    env_region,
-                    args.tupled_upvars_ty(),
-                    args.coroutine_captures_by_ref_ty(),
-                )
-            } else {
-                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;
-                // When we don't know the closure kind (and therefore also the closure's upvars,
-                // which are computed at the same time), we must delay the computation of the
-                // generator's upvars. We do this using the `AsyncFnKindHelper`, which as a trait
-                // goal functions similarly to the old `ClosureKind` predicate, and ensures that
-                // the goal kind <= the closure kind. As a projection `AsyncFnKindHelper::Upvars`
-                // will project to the right upvars for the generator, appending the inputs and
-                // coroutine upvars respecting the closure kind.
-                // N.B. No need to register a `AsyncFnKindHelper` goal here, it's already in `nested`.
-                let tupled_upvars_ty = Ty::new_projection(
+                    obligation.predicate.def_id,
+                    [self_ty, sig.tupled_inputs_ty],
+                ),
+                sym::CallMutFuture | sym::CallFuture => ty::AliasTy::new(
                     tcx,
-                    upvars_projection_def_id,
+                    obligation.predicate.def_id,
+                    [ty::GenericArg::from(self_ty), sig.tupled_inputs_ty.into(), env_region.into()],
+                ),
+                name => bug!("no such associated type: {name}"),
+            };
+
+            args.coroutine_closure_sig()
+                .rebind(ty::ProjectionPredicate { projection_ty, term: term.into() })
+        }
+        ty::FnDef(..) | ty::FnPtr(..) => {
+            let bound_sig = self_ty.fn_sig(tcx);
+            let sig = bound_sig.skip_binder();
+
+            let term = match item_name {
+                sym::CallOnceFuture | sym::CallMutFuture | sym::CallFuture => sig.output(),
+                sym::Output => {
+                    let future_trait_def_id = tcx.require_lang_item(LangItem::Future, None);
+                    let future_output_def_id = tcx
+                        .associated_items(future_trait_def_id)
+                        .filter_by_name_unhygienic(sym::Output)
+                        .next()
+                        .unwrap()
+                        .def_id;
+                    Ty::new_projection(tcx, future_output_def_id, [sig.output()])
+                }
+                name => bug!("no such associated type: {name}"),
+            };
+            let projection_ty = match item_name {
+                sym::CallOnceFuture | sym::Output => ty::AliasTy::new(
+                    tcx,
+                    obligation.predicate.def_id,
+                    [self_ty, Ty::new_tup(tcx, sig.inputs())],
+                ),
+                sym::CallMutFuture | sym::CallFuture => ty::AliasTy::new(
+                    tcx,
+                    obligation.predicate.def_id,
                     [
-                        ty::GenericArg::from(kind_ty),
-                        Ty::from_closure_kind(tcx, goal_kind).into(),
+                        ty::GenericArg::from(self_ty),
+                        Ty::new_tup(tcx, sig.inputs()).into(),
                         env_region.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,
-                )
-            }
+                ),
+                name => bug!("no such associated type: {name}"),
+            };
+
+            bound_sig.rebind(ty::ProjectionPredicate { projection_ty, term: term.into() })
         }
-        sym::Output => sig.return_ty,
-        name => bug!("no such associated type: {name}"),
-    };
-    let projection_ty = match item_name {
-        sym::CallOnceFuture | sym::Output => {
-            ty::AliasTy::new(tcx, obligation.predicate.def_id, [self_ty, sig.tupled_inputs_ty])
+        ty::Closure(_, args) => {
+            let args = args.as_closure();
+            let bound_sig = args.sig();
+            let sig = bound_sig.skip_binder();
+
+            let term = match item_name {
+                sym::CallOnceFuture | sym::CallMutFuture | sym::CallFuture => sig.output(),
+                sym::Output => {
+                    let future_trait_def_id = tcx.require_lang_item(LangItem::Future, None);
+                    let future_output_def_id = tcx
+                        .associated_items(future_trait_def_id)
+                        .filter_by_name_unhygienic(sym::Output)
+                        .next()
+                        .unwrap()
+                        .def_id;
+                    Ty::new_projection(tcx, future_output_def_id, [sig.output()])
+                }
+                name => bug!("no such associated type: {name}"),
+            };
+            let projection_ty = match item_name {
+                sym::CallOnceFuture | sym::Output => {
+                    ty::AliasTy::new(tcx, obligation.predicate.def_id, [self_ty, sig.inputs()[0]])
+                }
+                sym::CallMutFuture | sym::CallFuture => ty::AliasTy::new(
+                    tcx,
+                    obligation.predicate.def_id,
+                    [ty::GenericArg::from(self_ty), sig.inputs()[0].into(), env_region.into()],
+                ),
+                name => bug!("no such associated type: {name}"),
+            };
+
+            bound_sig.rebind(ty::ProjectionPredicate { projection_ty, term: term.into() })
         }
-        sym::CallMutFuture | sym::CallFuture => ty::AliasTy::new(
-            tcx,
-            obligation.predicate.def_id,
-            [ty::GenericArg::from(self_ty), sig.tupled_inputs_ty.into(), env_region.into()],
-        ),
-        name => bug!("no such associated type: {name}"),
+        _ => bug!("expected callable type for AsyncFn candidate"),
     };
 
-    confirm_param_env_candidate(
-        selcx,
-        obligation,
-        args.coroutine_closure_sig()
-            .rebind(ty::ProjectionPredicate { projection_ty, term: term.into() }),
-        true,
-    )
-    .with_addl_obligations(nested)
+    confirm_param_env_candidate(selcx, obligation, poly_cache_entry, true)
+        .with_addl_obligations(nested)
 }
 
 fn confirm_async_fn_kind_helper_candidate<'cx, 'tcx>(
diff --git a/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs b/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs
index 2258e796103..34dc8553714 100644
--- a/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs
+++ b/compiler/rustc_trait_selection/src/traits/select/candidate_assembly.rs
@@ -361,8 +361,18 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
                 }
                 candidates.vec.push(AsyncClosureCandidate);
             }
-            ty::Infer(ty::TyVar(_)) => {
-                candidates.ambiguous = true;
+            // Closures and fn pointers implement `AsyncFn*` if their return types
+            // implement `Future`, which is checked later.
+            ty::Closure(_, args) => {
+                if let Some(closure_kind) = args.as_closure().kind_ty().to_opt_closure_kind()
+                    && !closure_kind.extends(goal_kind)
+                {
+                    return;
+                }
+                candidates.vec.push(AsyncClosureCandidate);
+            }
+            ty::FnDef(..) | ty::FnPtr(..) => {
+                candidates.vec.push(AsyncClosureCandidate);
             }
             _ => {}
         }
diff --git a/compiler/rustc_trait_selection/src/traits/select/confirmation.rs b/compiler/rustc_trait_selection/src/traits/select/confirmation.rs
index c9d06b0f675..42845169549 100644
--- a/compiler/rustc_trait_selection/src/traits/select/confirmation.rs
+++ b/compiler/rustc_trait_selection/src/traits/select/confirmation.rs
@@ -883,40 +883,86 @@ impl<'cx, 'tcx> SelectionContext<'cx, 'tcx> {
         &mut self,
         obligation: &PolyTraitObligation<'tcx>,
     ) -> Result<Vec<PredicateObligation<'tcx>>, SelectionError<'tcx>> {
-        // Okay to skip binder because the args on closure types never
-        // touch bound regions, they just capture the in-scope
-        // type/region parameters.
+        let tcx = self.tcx();
         let self_ty = self.infcx.shallow_resolve(obligation.self_ty().skip_binder());
-        let ty::CoroutineClosure(closure_def_id, args) = *self_ty.kind() else {
-            bug!("async closure candidate for non-coroutine-closure {:?}", obligation);
-        };
 
-        let trait_ref = args.as_coroutine_closure().coroutine_closure_sig().map_bound(|sig| {
-            ty::TraitRef::new(
-                self.tcx(),
-                obligation.predicate.def_id(),
-                [self_ty, sig.tupled_inputs_ty],
-            )
-        });
+        let mut nested = vec![];
+        let (trait_ref, kind_ty) = match *self_ty.kind() {
+            ty::CoroutineClosure(_, args) => {
+                let args = args.as_coroutine_closure();
+                let trait_ref = args.coroutine_closure_sig().map_bound(|sig| {
+                    ty::TraitRef::new(
+                        self.tcx(),
+                        obligation.predicate.def_id(),
+                        [self_ty, sig.tupled_inputs_ty],
+                    )
+                });
+                (trait_ref, args.kind_ty())
+            }
+            ty::FnDef(..) | ty::FnPtr(..) => {
+                let sig = self_ty.fn_sig(tcx);
+                let trait_ref = sig.map_bound(|sig| {
+                    ty::TraitRef::new(
+                        self.tcx(),
+                        obligation.predicate.def_id(),
+                        [self_ty, Ty::new_tup(tcx, sig.inputs())],
+                    )
+                });
+                // We must additionally check that the return type impls `Future`.
+                let future_trait_def_id = tcx.require_lang_item(LangItem::Future, None);
+                nested.push(obligation.with(
+                    tcx,
+                    sig.map_bound(|sig| {
+                        ty::TraitRef::new(tcx, future_trait_def_id, [sig.output()])
+                    }),
+                ));
+                (trait_ref, Ty::from_closure_kind(tcx, ty::ClosureKind::Fn))
+            }
+            ty::Closure(_, args) => {
+                let sig = args.as_closure().sig();
+                let trait_ref = sig.map_bound(|sig| {
+                    ty::TraitRef::new(
+                        self.tcx(),
+                        obligation.predicate.def_id(),
+                        [self_ty, sig.inputs()[0]],
+                    )
+                });
+                // We must additionally check that the return type impls `Future`.
+                let future_trait_def_id = tcx.require_lang_item(LangItem::Future, None);
+                nested.push(obligation.with(
+                    tcx,
+                    sig.map_bound(|sig| {
+                        ty::TraitRef::new(tcx, future_trait_def_id, [sig.output()])
+                    }),
+                ));
+                (trait_ref, Ty::from_closure_kind(tcx, ty::ClosureKind::Fn))
+            }
+            _ => bug!("expected callable type for AsyncFn candidate"),
+        };
 
-        let mut nested = self.confirm_poly_trait_refs(obligation, trait_ref)?;
+        nested.extend(self.confirm_poly_trait_refs(obligation, trait_ref)?);
 
         let goal_kind =
             self.tcx().async_fn_trait_kind_from_def_id(obligation.predicate.def_id()).unwrap();
-        nested.push(obligation.with(
-            self.tcx(),
-            ty::TraitRef::from_lang_item(
-                self.tcx(),
-                LangItem::AsyncFnKindHelper,
-                obligation.cause.span,
-                [
-                    args.as_coroutine_closure().kind_ty(),
-                    Ty::from_closure_kind(self.tcx(), goal_kind),
-                ],
-            ),
-        ));
 
-        debug!(?closure_def_id, ?trait_ref, ?nested, "confirm closure candidate obligations");
+        // If we have not yet determiend the `ClosureKind` of the closure or coroutine-closure,
+        // then additionally register an `AsyncFnKindHelper` goal which will fail if the kind
+        // is constrained to an insufficient type later on.
+        if let Some(closure_kind) = self.infcx.shallow_resolve(kind_ty).to_opt_closure_kind() {
+            if !closure_kind.extends(goal_kind) {
+                return Err(SelectionError::Unimplemented);
+            }
+        } else {
+            nested.push(obligation.with(
+                self.tcx(),
+                ty::TraitRef::from_lang_item(
+                    self.tcx(),
+                    LangItem::AsyncFnKindHelper,
+                    obligation.cause.span,
+                    [kind_ty, Ty::from_closure_kind(self.tcx(), goal_kind)],
+                ),
+            ));
+        }
 
         Ok(nested)
     }
diff --git a/compiler/rustc_ty_utils/src/instance.rs b/compiler/rustc_ty_utils/src/instance.rs
index 9faad10dd14..bcc7c98ed69 100644
--- a/compiler/rustc_ty_utils/src/instance.rs
+++ b/compiler/rustc_ty_utils/src/instance.rs
@@ -306,6 +306,19 @@ fn resolve_associated_item<'tcx>(
                             Some(Instance::new(coroutine_closure_def_id, args))
                         }
                     }
+                    ty::Closure(closure_def_id, args) => {
+                        let trait_closure_kind = tcx.fn_trait_kind_from_def_id(trait_id).unwrap();
+                        Some(Instance::resolve_closure(
+                            tcx,
+                            closure_def_id,
+                            args,
+                            trait_closure_kind,
+                        ))
+                    }
+                    ty::FnDef(..) | ty::FnPtr(..) => Some(Instance {
+                        def: ty::InstanceDef::FnPtrShim(trait_item_id, rcvr_args.type_at(0)),
+                        args: rcvr_args,
+                    }),
                     _ => bug!(
                         "no built-in definition for `{trait_ref}::{}` for non-lending-closure type",
                         tcx.item_name(trait_item_id)
diff --git a/tests/ui/async-await/async-fn/simple.rs b/tests/ui/async-await/async-fn/simple.rs
index 99a5d56a309..172ede7098a 100644
--- a/tests/ui/async-await/async-fn/simple.rs
+++ b/tests/ui/async-await/async-fn/simple.rs
@@ -1,5 +1,5 @@
 // edition: 2021
-// check-pass
+// build-pass
 
 #![feature(async_fn_traits)]
 
diff --git a/tests/ui/did_you_mean/bad-assoc-ty.stderr b/tests/ui/did_you_mean/bad-assoc-ty.stderr
index eed01267224..3c474d19d1d 100644
--- a/tests/ui/did_you_mean/bad-assoc-ty.stderr
+++ b/tests/ui/did_you_mean/bad-assoc-ty.stderr
@@ -191,14 +191,7 @@ error[E0223]: ambiguous associated type
   --> $DIR/bad-assoc-ty.rs:33:10
    |
 LL | type H = Fn(u8) -> (u8)::Output;
-   |          ^^^^^^^^^^^^^^^^^^^^^^
-   |
-help: use fully-qualified syntax
-   |
-LL | type H = <(dyn Fn(u8) -> u8 + 'static) as AsyncFnOnce>::Output;
-   |          ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-LL | type H = <(dyn Fn(u8) -> u8 + 'static) as IntoFuture>::Output;
-   |          ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+   |          ^^^^^^^^^^^^^^^^^^^^^^ help: use fully-qualified syntax: `<(dyn Fn(u8) -> u8 + 'static) as IntoFuture>::Output`
 
 error[E0223]: ambiguous associated type
   --> $DIR/bad-assoc-ty.rs:39:19