Change FutureObj to use a *mut dyn Future fat pointer

Author: Nemo157

futures-core/src/future/future_obj.rs

@@ -1,4 +1,5 @@
 use core::{
+    mem,
     fmt,
     future::Future,
     marker::PhantomData,
@@ -13,22 +14,32 @@ use core::{
 /// take `dyn Trait` by value and `Box<dyn Trait>` is not available in no_std
 /// contexts.
 pub struct LocalFutureObj<'a, T> {
-    ptr: *mut (),
-    poll_fn: unsafe fn(*mut (), &mut Context<'_>) -> Poll<T>,
-    drop_fn: unsafe fn(*mut ()),
+    future: *mut (dyn Future<Output = T> + 'static),
+    drop_fn: unsafe fn(*mut (dyn Future<Output = T> + 'static)),
     _marker: PhantomData<&'a ()>,
 }
 
 impl<T> Unpin for LocalFutureObj<'_, T> {}
 
+unsafe fn remove_future_lifetime<'a, T>(ptr: *mut (dyn Future<Output = T> + 'a))
+    -> *mut (dyn Future<Output = T> + 'static)
+{
+    mem::transmute(ptr)
+}
+
+unsafe fn remove_drop_lifetime<'a, T>(ptr: unsafe fn (*mut (dyn Future<Output = T> + 'a)))
+    -> unsafe fn(*mut (dyn Future<Output = T> + 'static))
+{
+    mem::transmute(ptr)
+}
+
 impl<'a, T> LocalFutureObj<'a, T> {
     /// Create a `LocalFutureObj` from a custom trait object representation.
     #[inline]
     pub fn new<F: UnsafeFutureObj<'a, T> + 'a>(f: F) -> LocalFutureObj<'a, T> {
         LocalFutureObj {
-            ptr: f.into_raw(),
-            poll_fn: F::poll,
-            drop_fn: F::drop,
+            future: unsafe { remove_future_lifetime(f.into_raw()) },
+            drop_fn: unsafe { remove_drop_lifetime(F::drop) },
             _marker: PhantomData,
         }
     }
@@ -61,17 +72,17 @@ impl<T> Future for LocalFutureObj<'_, T> {
     type Output = T;
 
     #[inline]
-    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
         unsafe {
-            ((*self).poll_fn)((*self).ptr, cx)
+            Pin::new_unchecked(&mut *self.future).poll(cx)
         }
     }
 }
 
 impl<T> Drop for LocalFutureObj<'_, T> {
     fn drop(&mut self) {
         unsafe {
-            (self.drop_fn)(self.ptr)
+            (self.drop_fn)(self.future)
         }
     }
 }
@@ -119,115 +130,174 @@ impl<T> Future for FutureObj<'_, T> {
 }
 
 /// A custom implementation of a future trait object for `FutureObj`, providing
-/// a hand-rolled vtable.
+/// a vtable with drop support.
 ///
 /// This custom representation is typically used only in `no_std` contexts,
 /// where the default `Box`-based implementation is not available.
 ///
-/// The implementor must guarantee that it is safe to call `poll` repeatedly (in
-/// a non-concurrent fashion) with the result of `into_raw` until `drop` is
-/// called.
+/// # Safety
+///
+/// See the safety notes on individual methods for what guarantees an
+/// implementor must provide.
 pub unsafe trait UnsafeFutureObj<'a, T>: 'a {
-    /// Convert an owned instance into a (conceptually owned) void pointer.
-    fn into_raw(self) -> *mut ();
-
-    /// Poll the future represented by the given void pointer.
+    /// Convert an owned instance into a (conceptually owned) fat pointer.
     ///
     /// # Safety
     ///
-    /// The trait implementor must guarantee that it is safe to repeatedly call
-    /// `poll` with the result of `into_raw` until `drop` is called; such calls
-    /// are not, however, allowed to race with each other or with calls to
-    /// `drop`.
-    unsafe fn poll(ptr: *mut (), cx: &mut Context<'_>) -> Poll<T>;
+    /// ## Implementor
+    ///
+    /// The trait implementor must guarantee that it is safe to convert the
+    /// provided `*mut (dyn Future<Output = T> + 'a)` into a `Pin<&mut (dyn
+    /// Future<Output = T> + 'a)>` and call methods on it, non-reentrantly,
+    /// until `UnsafeFutureObj::drop` is called with it.
+    fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a);
 
-    /// Drops the future represented by the given void pointer.
+    /// Drops the future represented by the given fat pointer.
     ///
     /// # Safety
     ///
+    /// ## Implementor
+    ///
     /// The trait implementor must guarantee that it is safe to call this
-    /// function once per `into_raw` invocation; that call cannot race with
-    /// other calls to `drop` or `poll`.
-    unsafe fn drop(ptr: *mut ());
+    /// function once per `into_raw` invocation.
+    ///
+    /// ## Caller
+    ///
+    /// The caller must ensure:
+    ///
+    ///  * the pointer passed was obtained from an `into_raw` invocation from
+    ///    this same trait object
+    ///  * the pointer is not currently in use as a `Pin<&mut (dyn Future<Output
+    ///    = T> + 'a)>`
+    ///  * the pointer must not be used again after this function is called
+    unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a));
 }
 
 unsafe impl<'a, T, F> UnsafeFutureObj<'a, T> for &'a mut F
 where
     F: Future<Output = T> + Unpin + 'a
 {
-    fn into_raw(self) -> *mut () {
-        self as *mut F as *mut ()
+    fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) {
+        self as *mut dyn Future<Output = T>
     }
 
-    unsafe fn poll(ptr: *mut (), cx: &mut Context<'_>) -> Poll<T> {
-        let p: Pin<&mut F> = Pin::new_unchecked(&mut *(ptr as *mut F));
-        F::poll(p, cx)
+    unsafe fn drop(_ptr: *mut (dyn Future<Output = T> + 'a)) {}
+}
+
+unsafe impl<'a, T> UnsafeFutureObj<'a, T> for &'a mut (dyn Future<Output = T> + Unpin + 'a)
+{
+    fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) {
+        self as *mut dyn Future<Output = T>
     }
 
-    unsafe fn drop(_ptr: *mut ()) {}
+    unsafe fn drop(_ptr: *mut (dyn Future<Output = T> + 'a)) {}
 }
 
 unsafe impl<'a, T, F> UnsafeFutureObj<'a, T> for Pin<&'a mut F>
 where
     F: Future<Output = T> + 'a
 {
-    fn into_raw(mut self) -> *mut () {
-        let mut_ref: &mut F = unsafe { Pin::get_unchecked_mut(self.as_mut()) };
-        mut_ref as *mut F as *mut ()
+    fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) {
+        unsafe { self.get_unchecked_mut() as *mut dyn Future<Output = T> }
     }
 
-    unsafe fn poll(ptr: *mut (), cx: &mut Context<'_>) -> Poll<T> {
-        let future: Pin<&mut F> = Pin::new_unchecked(&mut *(ptr as *mut F));
-        F::poll(future, cx)
+    unsafe fn drop(_ptr: *mut (dyn Future<Output = T> + 'a)) {}
+}
+
+unsafe impl<'a, T> UnsafeFutureObj<'a, T> for Pin<&'a mut (dyn Future<Output = T> + 'a)>
+{
+    fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) {
+        unsafe { self.get_unchecked_mut() as *mut dyn Future<Output = T> }
     }
 
-    unsafe fn drop(_ptr: *mut ()) {}
+    unsafe fn drop(_ptr: *mut (dyn Future<Output = T> + 'a)) {}
 }
 
 #[cfg(feature = "alloc")]
 mod if_alloc {
     use super::*;
-    use core::mem;
     use alloc::boxed::Box;
 
     unsafe impl<'a, T, F> UnsafeFutureObj<'a, T> for Box<F>
         where F: Future<Output = T> + 'a
     {
-        fn into_raw(self) -> *mut () {
-            Box::into_raw(self) as *mut ()
+        fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) {
+            Box::into_raw(self)
         }
 
-        unsafe fn poll(ptr: *mut (), cx: &mut Context<'_>) -> Poll<T> {
-            let ptr = ptr as *mut F;
-            let pin: Pin<&mut F> = Pin::new_unchecked(&mut *ptr);
-            F::poll(pin, cx)
+        unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) {
+            drop(Box::from_raw(ptr as *mut F))
         }
+    }
 
-        unsafe fn drop(ptr: *mut ()) {
-            drop(Box::from_raw(ptr as *mut F))
+    unsafe impl<'a, T: 'a> UnsafeFutureObj<'a, T> for Box<dyn Future<Output = T> + 'a> {
+        fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) {
+            Box::into_raw(self)
+        }
+
+        unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) {
+            drop(Box::from_raw(ptr))
+        }
+    }
+
+    unsafe impl<'a, T: 'a> UnsafeFutureObj<'a, T> for Box<dyn Future<Output = T> + Send + 'a> {
+        fn into_raw(self) -> *mut (dyn Future<Output = T> + 'a) {
+            Box::into_raw(self)
+        }
+
+        unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) {
+            drop(Box::from_raw(ptr))
         }
     }
 
     unsafe impl<'a, T, F> UnsafeFutureObj<'a, T> for Pin<Box<F>>
     where
         F: Future<Output = T> + 'a
     {
-        fn into_raw(mut self) -> *mut () {
-            let mut_ref: &mut F = unsafe { Pin::get_unchecked_mut(self.as_mut()) };
-            let ptr = mut_ref as *mut F as *mut ();
-            mem::forget(self); // Don't drop the box
+        fn into_raw(mut self) -> *mut (dyn Future<Output = T> + 'a) {
+            let ptr = unsafe { self.as_mut().get_unchecked_mut() as *mut _ };
+            mem::forget(self);
+            ptr
+        }
+
+        unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) {
+            drop(Pin::from(Box::from_raw(ptr)))
+        }
+    }
+
+    unsafe impl<'a, T: 'a> UnsafeFutureObj<'a, T> for Pin<Box<dyn Future<Output = T> + 'a>> {
+        fn into_raw(mut self) -> *mut (dyn Future<Output = T> + 'a) {
+            let ptr = unsafe { self.as_mut().get_unchecked_mut() as *mut _ };
+            mem::forget(self);
+            ptr
+        }
+
+        unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) {
+            drop(Pin::from(Box::from_raw(ptr)))
+        }
+    }
+
+    unsafe impl<'a, T: 'a> UnsafeFutureObj<'a, T> for Pin<Box<dyn Future<Output = T> + Send + 'a>> {
+        fn into_raw(mut self) -> *mut (dyn Future<Output = T> + 'a) {
+            let ptr = unsafe { self.as_mut().get_unchecked_mut() as *mut _ };
+            mem::forget(self);
             ptr
         }
 
-        unsafe fn poll(ptr: *mut (), cx: &mut Context<'_>) -> Poll<T> {
-            let ptr = ptr as *mut F;
-            let pin: Pin<&mut F> = Pin::new_unchecked(&mut *ptr);
-            F::poll(pin, cx)
+        unsafe fn drop(ptr: *mut (dyn Future<Output = T> + 'a)) {
+            drop(Pin::from(Box::from_raw(ptr)))
+        }
+    }
+
+    impl<'a, F: Future<Output = ()> + Send + 'a> From<Box<F>> for FutureObj<'a, ()> {
+        fn from(boxed: Box<F>) -> Self {
+            FutureObj::new(boxed)
         }
+    }
 
-        unsafe fn drop(ptr: *mut ()) {
-            #[allow(clippy::cast_ptr_alignment)]
-            drop(Pin::from(Box::from_raw(ptr as *mut F)));
+    impl<'a> From<Box<dyn Future<Output = ()> + Send + 'a>> for FutureObj<'a, ()> {
+        fn from(boxed: Box<dyn Future<Output = ()> + Send + 'a>) -> Self {
+            FutureObj::new(boxed)
         }
     }
 
@@ -237,20 +307,32 @@ mod if_alloc {
         }
     }
 
-    impl<'a, F: Future<Output = ()> + Send + 'a> From<Box<F>> for FutureObj<'a, ()> {
-        fn from(boxed: Box<F>) -> Self {
+    impl<'a> From<Pin<Box<dyn Future<Output = ()> + Send + 'a>>> for FutureObj<'a, ()> {
+        fn from(boxed: Pin<Box<dyn Future<Output = ()> + Send + 'a>>) -> Self {
             FutureObj::new(boxed)
         }
     }
 
+    impl<'a, F: Future<Output = ()> + 'a> From<Box<F>> for LocalFutureObj<'a, ()> {
+        fn from(boxed: Box<F>) -> Self {
+            LocalFutureObj::new(boxed)
+        }
+    }
+
+    impl<'a> From<Box<dyn Future<Output = ()> + 'a>> for LocalFutureObj<'a, ()> {
+        fn from(boxed: Box<dyn Future<Output = ()> + 'a>) -> Self {
+            LocalFutureObj::new(boxed)
+        }
+    }
+
     impl<'a, F: Future<Output = ()> + 'a> From<Pin<Box<F>>> for LocalFutureObj<'a, ()> {
         fn from(boxed: Pin<Box<F>>) -> Self {
             LocalFutureObj::new(boxed)
         }
     }
 
-    impl<'a, F: Future<Output = ()> + 'a> From<Box<F>> for LocalFutureObj<'a, ()> {
-        fn from(boxed: Box<F>) -> Self {
+    impl<'a> From<Pin<Box<dyn Future<Output = ()> + 'a>>> for LocalFutureObj<'a, ()> {
+        fn from(boxed: Pin<Box<dyn Future<Output = ()> + 'a>>) -> Self {
             LocalFutureObj::new(boxed)
         }
     }

futures-util/src/future/mod.rs

@@ -493,8 +493,8 @@ pub trait FutureExt: Future {
 
     /// Wrap the future in a Box, pinning it.
     #[cfg(feature = "alloc")]
-    fn boxed(self) -> BoxFuture<'static, Self::Output>
-        where Self: Sized + Send + 'static
+    fn boxed<'a>(self) -> BoxFuture<'a, Self::Output>
+        where Self: Sized + Send + 'a
     {
         Box::pin(self)
     }

futures/tests/future_obj.rs

@@ -1,12 +1,12 @@
 #![feature(async_await)]
 
-use futures::future::{Future, FutureObj};
+use futures::future::{Future, FutureObj, FutureExt};
 use std::pin::Pin;
 use futures::task::{Context, Poll};
 
 #[test]
 fn dropping_does_not_segfault() {
-    FutureObj::new(Box::new(async { String::new() }));
+    FutureObj::new(async { String::new() }.boxed());
 }
 
 #[test]
@@ -29,7 +29,7 @@ fn dropping_drops_the_future() {
         }
     }
 
-    FutureObj::new(Box::new(Inc(&mut times_dropped)));
+    FutureObj::new(Inc(&mut times_dropped).boxed());
 
     assert_eq!(times_dropped, 1);
 }