make better async fn kind errors

7 files changed, 154 insertions, 37 deletions

diff --git a/compiler/rustc_trait_selection/messages.ftl b/compiler/rustc_trait_selection/messages.ftl
index 41db8059cbe..0dcba0e05f7 100644
--- a/compiler/rustc_trait_selection/messages.ftl
+++ b/compiler/rustc_trait_selection/messages.ftl
@@ -8,14 +8,16 @@ trait_selection_adjust_signature_remove_borrow = consider adjusting the signatur
         *[other] arguments
     }
 
-trait_selection_closure_fn_mut_label = closure is `FnMut` because it mutates the variable `{$place}` here
+trait_selection_async_closure_not_fn = async closure does not implement `{$kind}` because it captures state from its environment
 
-trait_selection_closure_fn_once_label = closure is `FnOnce` because it moves the variable `{$place}` out of its environment
+trait_selection_closure_fn_mut_label = closure is `{$trait_prefix}FnMut` because it mutates the variable `{$place}` here
 
-trait_selection_closure_kind_mismatch = expected a closure that implements the `{$expected}` trait, but this closure only implements `{$found}`
-    .label = this closure implements `{$found}`, not `{$expected}`
+trait_selection_closure_fn_once_label = closure is `{$trait_prefix}FnOnce` because it moves the variable `{$place}` out of its environment
 
-trait_selection_closure_kind_requirement = the requirement to implement `{$expected}` derives from here
+trait_selection_closure_kind_mismatch = expected a closure that implements the `{$trait_prefix}{$expected}` trait, but this closure only implements `{$trait_prefix}{$found}`
+    .label = this closure implements `{$trait_prefix}{$found}`, not `{$trait_prefix}{$expected}`
+
+trait_selection_closure_kind_requirement = the requirement to implement `{$trait_prefix}{$expected}` derives from here
 
 trait_selection_dump_vtable_entries = vtable entries for `{$trait_ref}`: {$entries}
 
diff --git a/compiler/rustc_trait_selection/src/errors.rs b/compiler/rustc_trait_selection/src/errors.rs
index 20cd573f46e..407fff03e15 100644
--- a/compiler/rustc_trait_selection/src/errors.rs
+++ b/compiler/rustc_trait_selection/src/errors.rs
@@ -135,6 +135,8 @@ pub struct ClosureKindMismatch {
     #[label(trait_selection_closure_kind_requirement)]
     pub cause_span: Span,
 
+    pub trait_prefix: &'static str,
+
     #[subdiagnostic]
     pub fn_once_label: Option<ClosureFnOnceLabel>,
 
@@ -157,3 +159,11 @@ pub struct ClosureFnMutLabel {
     pub span: Span,
     pub place: String,
 }
+
+#[derive(Diagnostic)]
+#[diag(trait_selection_async_closure_not_fn)]
+pub(crate) struct AsyncClosureNotFn {
+    #[primary_span]
+    pub span: Span,
+    pub kind: &'static str,
+}
diff --git a/compiler/rustc_trait_selection/src/traits/error_reporting/type_err_ctxt_ext.rs b/compiler/rustc_trait_selection/src/traits/error_reporting/type_err_ctxt_ext.rs
index 73effb33560..68b1a0d4e61 100644
--- a/compiler/rustc_trait_selection/src/traits/error_reporting/type_err_ctxt_ext.rs
+++ b/compiler/rustc_trait_selection/src/traits/error_reporting/type_err_ctxt_ext.rs
@@ -2,7 +2,9 @@
 
 use super::on_unimplemented::{AppendConstMessage, OnUnimplementedNote, TypeErrCtxtExt as _};
 use super::suggestions::{get_explanation_based_on_obligation, TypeErrCtxtExt as _};
-use crate::errors::{ClosureFnMutLabel, ClosureFnOnceLabel, ClosureKindMismatch};
+use crate::errors::{
+    AsyncClosureNotFn, ClosureFnMutLabel, ClosureFnOnceLabel, ClosureKindMismatch,
+};
 use crate::infer::error_reporting::{TyCategory, TypeAnnotationNeeded as ErrorCode};
 use crate::infer::type_variable::{TypeVariableOrigin, TypeVariableOriginKind};
 use crate::infer::InferCtxtExt as _;
@@ -959,34 +961,102 @@ impl<'tcx> TypeErrCtxtExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
     fn emit_specialized_closure_kind_error(
         &self,
         obligation: &PredicateObligation<'tcx>,
-        trait_ref: ty::PolyTraitRef<'tcx>,
+        mut trait_ref: ty::PolyTraitRef<'tcx>,
     ) -> Option<ErrorGuaranteed> {
-        let self_ty = trait_ref.self_ty().skip_binder();
-        if let ty::Closure(closure_def_id, closure_args) = *self_ty.kind()
-            && let Some(expected_kind) = self.tcx.fn_trait_kind_from_def_id(trait_ref.def_id())
-            && let Some(found_kind) = self.closure_kind(self_ty)
+        // If `AsyncFnKindHelper` is not implemented, that means that the closure kind
+        // doesn't extend the goal kind. This is worth reporting, but we can only do so
+        // if we actually know which closure this goal comes from, so look at the cause
+        // to see if we can extract that information.
+        if Some(trait_ref.def_id()) == self.tcx.lang_items().async_fn_kind_helper()
+            && let Some(found_kind) = trait_ref.skip_binder().args.type_at(0).to_opt_closure_kind()
+            && let Some(expected_kind) =
+                trait_ref.skip_binder().args.type_at(1).to_opt_closure_kind()
             && !found_kind.extends(expected_kind)
-            && let sig = closure_args.as_closure().sig()
-            && self.can_sub(
-                obligation.param_env,
-                trait_ref,
-                sig.map_bound(|sig| {
-                    ty::TraitRef::new(
-                        self.tcx,
-                        trait_ref.def_id(),
-                        [trait_ref.self_ty().skip_binder(), sig.inputs()[0]],
-                    )
-                }),
-            )
         {
-            let mut err =
-                self.report_closure_error(&obligation, closure_def_id, found_kind, expected_kind);
-            self.note_obligation_cause(&mut err, &obligation);
-            self.point_at_returns_when_relevant(&mut err, &obligation);
-            Some(err.emit())
-        } else {
-            None
+            if let Some((_, Some(parent))) = obligation.cause.code().parent() {
+                // If we have a derived obligation, then the parent will be a `AsyncFn*` goal.
+                trait_ref = parent.to_poly_trait_ref();
+            } else if let &ObligationCauseCode::FunctionArgumentObligation { arg_hir_id, .. } =
+                obligation.cause.code()
+                && let Some(typeck_results) = &self.typeck_results
+                && let ty::Closure(closure_def_id, _) | ty::CoroutineClosure(closure_def_id, _) =
+                    *typeck_results.node_type(arg_hir_id).kind()
+            {
+                // Otherwise, extract the closure kind from the obligation.
+                let mut err = self.report_closure_error(
+                    &obligation,
+                    closure_def_id,
+                    found_kind,
+                    expected_kind,
+                    "async ",
+                );
+                self.note_obligation_cause(&mut err, &obligation);
+                self.point_at_returns_when_relevant(&mut err, &obligation);
+                return Some(err.emit());
+            }
+        }
+
+        let self_ty = trait_ref.self_ty().skip_binder();
+
+        if let Some(expected_kind) = self.tcx.fn_trait_kind_from_def_id(trait_ref.def_id()) {
+            let (closure_def_id, found_args, by_ref_captures) = match *self_ty.kind() {
+                ty::Closure(def_id, args) => {
+                    (def_id, args.as_closure().sig().map_bound(|sig| sig.inputs()[0]), None)
+                }
+                ty::CoroutineClosure(def_id, args) => (
+                    def_id,
+                    args.as_coroutine_closure()
+                        .coroutine_closure_sig()
+                        .map_bound(|sig| sig.tupled_inputs_ty),
+                    Some(args.as_coroutine_closure().coroutine_captures_by_ref_ty()),
+                ),
+                _ => return None,
+            };
+
+            let expected_args = trait_ref.map_bound(|trait_ref| trait_ref.args.type_at(1));
+
+            // Verify that the arguments are compatible. If the signature is
+            // mismatched, then we have a totally different error to report.
+            if self.enter_forall(found_args, |found_args| {
+                self.enter_forall(expected_args, |expected_args| {
+                    !self.can_sub(obligation.param_env, expected_args, found_args)
+                })
+            }) {
+                return None;
+            }
+
+            if let Some(found_kind) = self.closure_kind(self_ty)
+                && !found_kind.extends(expected_kind)
+            {
+                let mut err = self.report_closure_error(
+                    &obligation,
+                    closure_def_id,
+                    found_kind,
+                    expected_kind,
+                    "",
+                );
+                self.note_obligation_cause(&mut err, &obligation);
+                self.point_at_returns_when_relevant(&mut err, &obligation);
+                return Some(err.emit());
+            }
+
+            // If the closure has captures, then perhaps the reason that the trait
+            // is unimplemented is because async closures don't implement `Fn`/`FnMut`
+            // if they have captures.
+            if let Some(by_ref_captures) = by_ref_captures
+                && let ty::FnPtr(sig) = by_ref_captures.kind()
+                && !sig.skip_binder().output().is_unit()
+            {
+                let mut err = self.tcx.dcx().create_err(AsyncClosureNotFn {
+                    span: self.tcx.def_span(closure_def_id),
+                    kind: expected_kind.as_str(),
+                });
+                self.note_obligation_cause(&mut err, &obligation);
+                self.point_at_returns_when_relevant(&mut err, &obligation);
+                return Some(err.emit());
+            }
         }
+        None
     }
 
     fn fn_arg_obligation(
@@ -1493,6 +1563,7 @@ pub(super) trait InferCtxtPrivExt<'tcx> {
         closure_def_id: DefId,
         found_kind: ty::ClosureKind,
         kind: ty::ClosureKind,
+        trait_prefix: &'static str,
     ) -> DiagnosticBuilder<'tcx>;
 
     fn report_cyclic_signature_error(
@@ -3376,6 +3447,7 @@ impl<'tcx> InferCtxtPrivExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
         closure_def_id: DefId,
         found_kind: ty::ClosureKind,
         kind: ty::ClosureKind,
+        trait_prefix: &'static str,
     ) -> DiagnosticBuilder<'tcx> {
         let closure_span = self.tcx.def_span(closure_def_id);
 
@@ -3384,6 +3456,7 @@ impl<'tcx> InferCtxtPrivExt<'tcx> for TypeErrCtxt<'_, 'tcx> {
             expected: kind,
             found: found_kind,
             cause_span: obligation.cause.span,
+            trait_prefix,
             fn_once_label: None,
             fn_mut_label: None,
         };
diff --git a/tests/ui/async-await/async-closures/not-fn.rs b/tests/ui/async-await/async-closures/not-fn.rs
new file mode 100644
index 00000000000..4505e6243e9
--- /dev/null
+++ b/tests/ui/async-await/async-closures/not-fn.rs
@@ -0,0 +1,15 @@
+// edition:2021
+
+// FIXME(async_closures): This needs a better error message!
+
+#![feature(async_closure)]
+
+fn main() {
+    fn needs_fn<T>(_: impl FnMut() -> T) {}
+
+    let mut x = 1;
+    needs_fn(async || {
+        //~^ ERROR  async closure does not implement `FnMut` because it captures state from its environment
+        x += 1;
+    });
+}
diff --git a/tests/ui/async-await/async-closures/not-fn.stderr b/tests/ui/async-await/async-closures/not-fn.stderr
new file mode 100644
index 00000000000..9c40613599a
--- /dev/null
+++ b/tests/ui/async-await/async-closures/not-fn.stderr
@@ -0,0 +1,16 @@
+error: async closure does not implement `FnMut` because it captures state from its environment
+  --> $DIR/not-fn.rs:11:14
+   |
+LL |     needs_fn(async || {
+   |     -------- ^^^^^^^^
+   |     |
+   |     required by a bound introduced by this call
+   |
+note: required by a bound in `needs_fn`
+  --> $DIR/not-fn.rs:8:28
+   |
+LL |     fn needs_fn<T>(_: impl FnMut() -> T) {}
+   |                            ^^^^^^^^^^^^ required by this bound in `needs_fn`
+
+error: aborting due to 1 previous error
+
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 f86cee3e070..1e372deb984 100644
--- a/tests/ui/async-await/async-closures/wrong-fn-kind.rs
+++ b/tests/ui/async-await/async-closures/wrong-fn-kind.rs
@@ -9,8 +9,7 @@ fn main() {
 
     let mut x = 1;
     needs_async_fn(async || {
-        //~^ ERROR i16: ops::async_function::internal_implementation_detail::AsyncFnKindHelper<i8>
-        // FIXME: Should say "closure is `async FnMut` but it needs `async Fn`" or sth.
+        //~^ ERROR expected a closure that implements the `async Fn` trait, but this closure only implements `async FnMut`
         x += 1;
     });
 }
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 4ef8484cc34..34a6b3a485a 100644
--- a/tests/ui/async-await/async-closures/wrong-fn-kind.stderr
+++ b/tests/ui/async-await/async-closures/wrong-fn-kind.stderr
@@ -1,15 +1,17 @@
-error[E0277]: the trait bound `i16: ops::async_function::internal_implementation_detail::AsyncFnKindHelper<i8>` is not satisfied
+error[E0525]: expected a closure that implements the `async Fn` trait, but this closure only implements `async FnMut`
   --> $DIR/wrong-fn-kind.rs:11:20
    |
 LL |       needs_async_fn(async || {
-   |  _____--------------_^
+   |       -------------- -^^^^^^^
+   |       |              |
+   |  _____|______________this closure implements `async FnMut`, not `async Fn`
    | |     |
    | |     required by a bound introduced by this call
 LL | |
-LL | |         // FIXME: Should say "closure is `async FnMut` but it needs `async Fn`" or sth.
 LL | |         x += 1;
+   | |         - closure is `async FnMut` because it mutates the variable `x` here
 LL | |     });
-   | |_____^ the trait `ops::async_function::internal_implementation_detail::AsyncFnKindHelper<i8>` is not implemented for `i16`
+   | |_____- the requirement to implement `async Fn` derives from here
    |
 note: required by a bound in `needs_async_fn`
   --> $DIR/wrong-fn-kind.rs:8:31
@@ -19,4 +21,4 @@ LL |     fn needs_async_fn(_: impl async Fn()) {}
 
 error: aborting due to 1 previous error
 
-For more information about this error, try `rustc --explain E0277`.
+For more information about this error, try `rustc --explain E0525`.