From 0dd40786b555c04afa52b9d0c789a29dbd4e3dd2 Mon Sep 17 00:00:00 2001
From: Michael Goulet <michael@errs.io>
Date: Wed, 31 Jan 2024 18:22:16 +0000
Subject: [PATCH] Harmonize blanket implementations for AsyncFn* traits

---
 library/alloc/src/boxed.rs             | 29 ++++++++++++++
 library/alloc/src/lib.rs               |  1 +
 library/core/src/ops/async_function.rs | 53 ++++++++++++++++++--------
 3 files changed, 68 insertions(+), 15 deletions(-)

diff --git a/library/alloc/src/boxed.rs b/library/alloc/src/boxed.rs
index 953041b8c20..92600b8e5bd 100644
--- a/library/alloc/src/boxed.rs
+++ b/library/alloc/src/boxed.rs
@@ -159,6 +159,7 @@ use core::iter::FusedIterator;
 use core::marker::Tuple;
 use core::marker::Unsize;
 use core::mem::{self, SizedTypeProperties};
+use core::ops::{AsyncFn, AsyncFnMut, AsyncFnOnce};
 use core::ops::{
     CoerceUnsized, Coroutine, CoroutineState, Deref, DerefMut, DispatchFromDyn, Receiver,
 };
@@ -2030,6 +2031,34 @@ impl<Args: Tuple, F: Fn<Args> + ?Sized, A: Allocator> Fn<Args> for Box<F, A> {
     }
 }
 
+#[unstable(feature = "async_fn_traits", issue = "none")]
+impl<Args: Tuple, F: AsyncFnOnce<Args> + ?Sized, A: Allocator> AsyncFnOnce<Args> for Box<F, A> {
+    type Output = F::Output;
+    type CallOnceFuture = F::CallOnceFuture;
+
+    extern "rust-call" fn async_call_once(self, args: Args) -> Self::CallOnceFuture {
+        F::async_call_once(*self, args)
+    }
+}
+
+#[unstable(feature = "async_fn_traits", issue = "none")]
+impl<Args: Tuple, F: AsyncFnMut<Args> + ?Sized, A: Allocator> AsyncFnMut<Args> for Box<F, A> {
+    type CallMutFuture<'a> = F::CallMutFuture<'a> where Self: 'a;
+
+    extern "rust-call" fn async_call_mut(&mut self, args: Args) -> Self::CallMutFuture<'_> {
+        F::async_call_mut(self, args)
+    }
+}
+
+#[unstable(feature = "async_fn_traits", issue = "none")]
+impl<Args: Tuple, F: AsyncFn<Args> + ?Sized, A: Allocator> AsyncFn<Args> for Box<F, A> {
+    type CallFuture<'a> = F::CallFuture<'a> where Self: 'a;
+
+    extern "rust-call" fn async_call(&self, args: Args) -> Self::CallFuture<'_> {
+        F::async_call(self, args)
+    }
+}
+
 #[unstable(feature = "coerce_unsized", issue = "18598")]
 impl<T: ?Sized + Unsize<U>, U: ?Sized, A: Allocator> CoerceUnsized<Box<U, A>> for Box<T, A> {}
 
diff --git a/library/alloc/src/lib.rs b/library/alloc/src/lib.rs
index 96d43e11dc6..3341b564d1f 100644
--- a/library/alloc/src/lib.rs
+++ b/library/alloc/src/lib.rs
@@ -106,6 +106,7 @@
 #![feature(array_windows)]
 #![feature(ascii_char)]
 #![feature(assert_matches)]
+#![feature(async_fn_traits)]
 #![feature(async_iterator)]
 #![feature(coerce_unsized)]
 #![feature(const_align_of_val)]
diff --git a/library/core/src/ops/async_function.rs b/library/core/src/ops/async_function.rs
index efbe9d164c3..19b1220f05e 100644
--- a/library/core/src/ops/async_function.rs
+++ b/library/core/src/ops/async_function.rs
@@ -65,44 +65,67 @@ pub trait AsyncFnOnce<Args: Tuple> {
 
 mod impls {
     use super::{AsyncFn, AsyncFnMut, AsyncFnOnce};
-    use crate::future::Future;
     use crate::marker::Tuple;
 
     #[unstable(feature = "async_fn_traits", issue = "none")]
-    impl<F: Fn<A>, A: Tuple> AsyncFn<A> for F
+    impl<A: Tuple, F: ?Sized> AsyncFn<A> for &F
     where
-        <F as FnOnce<A>>::Output: Future,
+        F: AsyncFn<A>,
     {
-        type CallFuture<'a> = <F as FnOnce<A>>::Output where Self: 'a;
+        type CallFuture<'a> = F::CallFuture<'a> where Self: 'a;
 
         extern "rust-call" fn async_call(&self, args: A) -> Self::CallFuture<'_> {
-            self.call(args)
+            F::async_call(*self, args)
         }
     }
 
     #[unstable(feature = "async_fn_traits", issue = "none")]
-    impl<F: FnMut<A>, A: Tuple> AsyncFnMut<A> for F
+    impl<A: Tuple, F: ?Sized> AsyncFnMut<A> for &F
     where
-        <F as FnOnce<A>>::Output: Future,
+        F: AsyncFn<A>,
     {
-        type CallMutFuture<'a> = <F as FnOnce<A>>::Output where Self: 'a;
+        type CallMutFuture<'a> = F::CallFuture<'a> where Self: 'a;
 
         extern "rust-call" fn async_call_mut(&mut self, args: A) -> Self::CallMutFuture<'_> {
-            self.call_mut(args)
+            F::async_call(*self, args)
         }
     }
 
     #[unstable(feature = "async_fn_traits", issue = "none")]
-    impl<F: FnOnce<A>, A: Tuple> AsyncFnOnce<A> for F
+    impl<'a, A: Tuple, F: ?Sized> AsyncFnOnce<A> for &'a F
     where
-        <F as FnOnce<A>>::Output: Future,
+        F: AsyncFn<A>,
     {
-        type CallOnceFuture = <F as FnOnce<A>>::Output;
-
-        type Output = <<F as FnOnce<A>>::Output as Future>::Output;
+        type Output = F::Output;
+        type CallOnceFuture = F::CallFuture<'a>;
 
         extern "rust-call" fn async_call_once(self, args: A) -> Self::CallOnceFuture {
-            self.call_once(args)
+            F::async_call(self, args)
+        }
+    }
+
+    #[unstable(feature = "async_fn_traits", issue = "none")]
+    impl<A: Tuple, F: ?Sized> AsyncFnMut<A> for &mut F
+    where
+        F: AsyncFnMut<A>,
+    {
+        type CallMutFuture<'a> = F::CallMutFuture<'a> where Self: 'a;
+
+        extern "rust-call" fn async_call_mut(&mut self, args: A) -> Self::CallMutFuture<'_> {
+            F::async_call_mut(*self, args)
+        }
+    }
+
+    #[unstable(feature = "async_fn_traits", issue = "none")]
+    impl<'a, A: Tuple, F: ?Sized> AsyncFnOnce<A> for &'a mut F
+    where
+        F: AsyncFnMut<A>,
+    {
+        type Output = F::Output;
+        type CallOnceFuture = F::CallMutFuture<'a>;
+
+        extern "rust-call" fn async_call_once(self, args: A) -> Self::CallOnceFuture {
+            F::async_call_mut(self, args)
         }
     }
 }