From 762febdaf33667995578ec8183a2379e28bb6b40 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Tue, 20 Feb 2024 15:43:29 +0000
Subject: [PATCH] Fix stray trait mismatch in resolve_associated_item for
 AsyncFn

---
 compiler/rustc_ty_utils/src/instance.rs       | 18 +++--------------
 .../async-fn/auxiliary/block-on.rs            | 20 +++++++++++++++++++
 tests/ui/async-await/async-fn/simple.rs       |  7 ++++++-
 3 files changed, 29 insertions(+), 16 deletions(-)
 create mode 100644 tests/ui/async-await/async-fn/auxiliary/block-on.rs

diff --git a/compiler/rustc_ty_utils/src/instance.rs b/compiler/rustc_ty_utils/src/instance.rs
index 5fc93d666ab16..3e3bccce47fe4 100644
--- a/compiler/rustc_ty_utils/src/instance.rs
+++ b/compiler/rustc_ty_utils/src/instance.rs
@@ -246,7 +246,7 @@ fn resolve_associated_item<'tcx>(
                         span: tcx.def_span(trait_item_id),
                     })
                 }
-            } else if tcx.fn_trait_kind_from_def_id(trait_ref.def_id).is_some() {
+            } else if let Some(target_kind) = tcx.fn_trait_kind_from_def_id(trait_ref.def_id) {
                 // FIXME: This doesn't check for malformed libcore that defines, e.g.,
                 // `trait Fn { fn call_once(&self) { .. } }`. This is mostly for extension
                 // methods.
@@ -265,13 +265,7 @@ fn resolve_associated_item<'tcx>(
                 }
                 match *rcvr_args.type_at(0).kind() {
                     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,
-                        ))
+                        Some(Instance::resolve_closure(tcx, closure_def_id, args, target_kind))
                     }
                     ty::FnDef(..) | ty::FnPtr(..) => Some(Instance {
                         def: ty::InstanceDef::FnPtrShim(trait_item_id, rcvr_args.type_at(0)),
@@ -324,13 +318,7 @@ fn resolve_associated_item<'tcx>(
                         }
                     }
                     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,
-                        ))
+                        Some(Instance::resolve_closure(tcx, closure_def_id, args, target_kind))
                     }
                     ty::FnDef(..) | ty::FnPtr(..) => Some(Instance {
                         def: ty::InstanceDef::FnPtrShim(trait_item_id, rcvr_args.type_at(0)),
diff --git a/tests/ui/async-await/async-fn/auxiliary/block-on.rs b/tests/ui/async-await/async-fn/auxiliary/block-on.rs
new file mode 100644
index 0000000000000..dcb710fc97c97
--- /dev/null
+++ b/tests/ui/async-await/async-fn/auxiliary/block-on.rs
@@ -0,0 +1,20 @@
+//@ edition: 2021
+
+#![feature(async_closure, noop_waker)]
+
+use std::future::Future;
+use std::pin::pin;
+use std::task::*;
+
+pub fn block_on<T>(fut: impl Future<Output = T>) -> T {
+    let mut fut = pin!(fut);
+    // Poll loop, just to test the future...
+    let ctx = &mut Context::from_waker(Waker::noop());
+
+    loop {
+        match unsafe { fut.as_mut().poll(ctx) } {
+            Poll::Pending => {}
+            Poll::Ready(t) => break t,
+        }
+    }
+}
diff --git a/tests/ui/async-await/async-fn/simple.rs b/tests/ui/async-await/async-fn/simple.rs
index e2a183a8c0ba0..21972ba5aefad 100644
--- a/tests/ui/async-await/async-fn/simple.rs
+++ b/tests/ui/async-await/async-fn/simple.rs
@@ -1,8 +1,11 @@
+//@ aux-build:block-on.rs
 //@ edition: 2021
 //@ build-pass
 
 #![feature(async_fn_traits)]
 
+extern crate block_on;
+
 use std::ops::AsyncFn;
 
 async fn foo() {}
@@ -12,5 +15,7 @@ async fn call_asyncly(f: impl AsyncFn(i32) -> i32) -> i32 {
 }
 
 fn main() {
-    let fut = call_asyncly(|x| async move { x + 1 });
+    block_on::block_on(async {
+        call_asyncly(|x| async move { x + 1 }).await;
+    });
 }