Harmonize blanket implementations for AsyncFn* traits

3 files changed, 67 insertions, 14 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::Output;
+        type CallOnceFuture = F::CallFuture<'a>;
 
-        type Output = <<F as FnOnce<A>>::Output as Future>::Output;
+        extern "rust-call" fn async_call_once(self, args: A) -> Self::CallOnceFuture {
+            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 {
-            self.call_once(args)
+            F::async_call_mut(self, args)
         }
     }
 }