Support self-contained strategies

[loichyan]

Resolves #19

Well, after ~2 days of research, I found that this feature is far more complicated than I initially thought. Since I have to prepare for an important exam, I must postpone the implementation (in the worst case, until the end of this year 😢). Nevertheless, I'm going to write about the progress I've already made.

Let's start with the goals we expected:

1. Provide self-contained containers so that each async method can return an allocated trait object by some strategy instead of a constructor. Currently, the only such container is Box. But using Box would lose the advantage that dynify brings. Therefore, we should provide containers that allocate objects on the stack.
2. Keep compatible with the current codebase of dynify¹ and stable Rust. If possible, provide a new feature flag that tells dynify to take advantage of nightly features to extend itself.

For (1), stackfuture provides a solution for dyn Future. However, I want dynify to work with any arbitrary traits (or at least ones in std). I found a nearly perfect approach for this:

use core::alloc::Layout;
use core::mem::MaybeUninit;
use core::ops::{Deref, DerefMut};
use core::pin::Pin;
use core::ptr::NonNull;
use core::task::{Context, Poll};

// Pointer that should never be dereferenced directly.
pub type VoidPtr = NonNull<Void>;
pub enum Void {}

// A weaker version of std's `Pointee` trait.
// See <https://doc.rust-lang.org/std/ptr/trait.Pointee.html>.
pub trait Pointee {
    type Metadata: Copy;
    // This is required since we cannot form a raw pointer to a dyn object
    // without the nightly `ptr_metadata` feature.
    unsafe fn drop_in_place(data: VoidPtr, metadata: Self::Metadata);
}

// A blanket trait that is used to mark sized `Pointee` types.
// For unsized types, we need to implement `Pointee` manually.
pub trait PointeeSized: Sized + Pointee<Metadata = ()> {}
impl<T: Sized> Pointee for T {
    // For sized types, no metadata is needed.
    type Metadata = ();
    unsafe fn drop_in_place(data: VoidPtr, _metadata: Self::Metadata) {
        data.cast::<T>().drop_in_place();
    }
}
impl<T: Sized> PointeeSized for T {}

// Since obtaining the metadata of a pointer to a dyn object requires nightly
// Rust, and a self-contained container requires the metadata to work properly,
// we need this trait to obtain the metadata from the corresponding sized type
// that is intended to be coerced into a dyn object.
pub trait CoerceInto<T: ?Sized + Pointee>: Pointee {
    fn coerce(data: VoidPtr, metadata: Self::Metadata) -> T::Metadata;
}

// Self-to-Self coercion should always be possible.
impl<T: ?Sized + Pointee> CoerceInto<T> for T {
    fn coerce(_data: VoidPtr, metadata: Self::Metadata) -> <T as Pointee>::Metadata {
        metadata
    }
}

// Metadata and vtable for `Future`s.
pub type FutureMetadata = &'static FutureVtable;
pub struct FutureVtable {
    // Add `Output` as a generic would bring unnecessary constraints.
    poll: unsafe fn(this: VoidPtr, cx: &mut Context, output: VoidPtr),
    drop_in_place: unsafe fn(this: VoidPtr),
}
impl FutureVtable {
    const fn new<T: Future>() -> &'static Self {
        &Self {
            poll: |this, cx, output| unsafe {
                let this = Pin::new_unchecked(this.cast::<T>().as_mut());
                let output = output.cast::<MaybeUninit<Poll<T::Output>>>().as_mut();
                output.write(this.poll(cx));
            },
            drop_in_place: |this| unsafe { this.cast::<T>().drop_in_place() },
        }
    }
}

// Add support for `Future`s.
// Implementations for other traits can be done in a similar way.
impl<'a, Output> Pointee for dyn 'a + Future<Output = Output> {
    type Metadata = FutureMetadata;
    unsafe fn drop_in_place(data: VoidPtr, metadata: Self::Metadata) {
        (metadata.drop_in_place)(data);
    }
}
impl<'a, T: 'a + Future> CoerceInto<dyn 'a + Future<Output = T::Output>> for T {
    fn coerce(_data: VoidPtr, _metadata: Self::Metadata) -> FutureMetadata {
        FutureVtable::new::<T>()
    }
}

#[repr(C)] // ensure `data` is at the beginning
pub struct Stack<T: ?Sized + Pointee> {
    data: MaybeUninit<[u8; 32]>,
    metadata: T::Metadata,
}
impl<T: ?Sized + Pointee> Stack<T> {
    // Initialize the given constructor and perform the intended coercion, which
    // may be a coercion from the sized `U` into the unsized dyn trait `T`, or
    // simply a Self-to-Self noop coercion.
    // Also note that here and the following we will use a simplified version of
    // the `Construct` trait for demonstration purpose.
    fn emplace_as<U>(
        constructor: impl FnOnce(VoidPtr) -> NonNull<U>,
        layout: Layout,
    ) -> Option<Self>
    where
        U: PointeeSized + CoerceInto<T>,
    {
        if layout.align() > std::mem::align_of::<Self>() || layout.size() > 32 {
            return None;
        }

        unsafe {
            let mut data = MaybeUninit::<[u8; 32]>::uninit();
            let data_ptr = NonNull::from(&mut data).cast();
            constructor(data_ptr);
            let metadata = U::coerce(NonNull::from(&data).cast(), ());
            Some(Self { data, metadata })
        }
    }
}

// This trait implementation is required by our `Emplace` trait. We cannot
// implement this for unsized types since there's no way to form a fat pointer
// without `ptr_metadata`.
// Though we can relax the requirement of `Emplace::Ptr`, this would be a
// breaking change.
impl<T: PointeeSized> Deref for Stack<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        unsafe { NonNull::from(&self.data).cast().as_ref() }
    }
}
impl<T: PointeeSized> DerefMut for Stack<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { NonNull::from(&mut self.data).cast().as_mut() }
    }
}

// `drop()` is handled by the `Pointee`.
impl<T: ?Sized + Pointee> Drop for Stack<T> {
    fn drop(&mut self) {
        unsafe {
            T::drop_in_place(NonNull::from(&mut self.data).cast(), self.metadata);
        }
    }
}

// Inherit marker traits.
impl<T: ?Sized + Pointee + Unpin> Unpin for Stack<T> {}
unsafe impl<T: ?Sized + Pointee + Send> Send for Stack<T> {}
unsafe impl<T: ?Sized + Pointee + Sync> Sync for Stack<T> {}

// For each supported trait, we must implement it for our containers as well.
impl<'a, T> Future for Stack<T>
where
    T: 'a + ?Sized + Future + CoerceInto<dyn 'a + Future<Output = T::Output>>,
{
    type Output = T::Output;
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        let mut output = MaybeUninit::uninit();
        unsafe {
            let this = self.get_unchecked_mut();
            let data = NonNull::from(&mut this.data).cast();
            let metadata = T::coerce(data, this.metadata);
            (metadata.poll)(data, cx, NonNull::from(&mut output).cast());
            output.assume_init()
        }
    }
}

// This defines the main entrypoint of all supported strategies.
pub trait Strategy<T: ?Sized + Pointee>: Default {
    type Output;
    fn apply<U>(self, input: impl FnOnce(VoidPtr) -> NonNull<U>, layout: Layout) -> Self::Output
    where
        U: PointeeSized + CoerceInto<T>;
}

// `Strategy` implementation for `Stack`.
// We can provide const generics to customize the stack size and alignment.
#[derive(Default)]
pub struct Stacked;
impl<T: ?Sized + Pointee> Strategy<T> for Stacked {
    type Output = Option<Stack<T>>;
    fn apply<U>(self, input: impl FnOnce(VoidPtr) -> NonNull<U>, layout: Layout) -> Self::Output
    where
        U: PointeeSized + CoerceInto<T>,
    {
        Stack::emplace_as(input, layout)
    }
}

Below is an example of macro expansion:

// With: #[dynify(strategy = "Stacked")]
pub trait Async {
    async fn test(&self) -> String;
}
// Output:
pub trait DynAsync {
    fn test<'this, 'dynify>(
        &'this self,
    ) -> <Stacked as Strategy<dyn 'dynify + Future<Output = String>>>::Output
    where
        'this: 'dynify;
}
impl<T: Async> DynAsync for T {
    fn test<'this, 'dynify>(
        &'this self,
    ) -> <Stacked as Strategy<dyn 'dynify + Future<Output = String>>>::Output
    where
        'this: 'dynify,
    {
        Strategy::apply(<Stacked as Default>::default(), || T::test(self))
    }
}

Although I found several other possible solutions², this one is the closest to perfect:

1. It is completely compatible with the current codebase of dynify.

2. The extra traits to be introduced are adapted from std's nightly features:

   • Pointee corresponds to Pointee.
   • CoerceInto corresponds to Unsize (though its name is closer to CoerceUnsized).

   This makes it convenient to smoothly migrate to nightly Rust (as discussed in the following).

3. In theory, we can support any arbitrary traits with these two traits (though I don't plan to export them yet since Future and Any are the most useful ones).

For (2), we can provide a new feature flag to enable support for arbitrary traits on nightly Rust. With that flag enabled, we're able to drop the need for manually crafted vtables and relax the trait bounds for some implementations:

// We can simply use std's `Pointee` trait.
pub use core::ptr::Pointee;

// `PointeeSized` is no longer needed.
// pub trait PointeeSized: Sized + Pointee<Metadata = ()> {}

// Technically, we can also drop this trait, but it's still required to keep
// compatible with dynify on stable Rust.
pub trait CoerceInto<T: ?Sized + Pointee>: Pointee {
    fn coerce(data: VoidPtr, metadata: Self::Metadata) -> T::Metadata;
}

// Additionally, we can implement `CoerceInto` for any type with `Unsize`, but
// this would require another nightly feature, which I want to avoid if possible.
//
// impl<T, U> CoerceInto<U> for T
// where
//     T: ?Sized + core::marker::Unsize<U>,
//     U: ?Sized,
// {
//     fn coerce(data: VoidPtr, metadata: Self::Metadata) -> <U as Pointee>::Metadata {
//         let ptr = NonNull::<T>::from_raw_parts(data, metadata);
//         (ptr as NonNull<U>).to_raw_parts().1
//     }
// }

// Self-to-Self coercion should always be possible.
impl<T: ?Sized + Pointee> CoerceInto<T> for T {
    fn coerce(_data: VoidPtr, metadata: Self::Metadata) -> <T as Pointee>::Metadata {
        metadata
    }
}

impl<'a, T: 'a + Future> CoerceInto<dyn 'a + Future<Output = T::Output>> for T {
    fn coerce(
        data: VoidPtr,
        _metadata: Self::Metadata,
    ) -> <dyn 'a + Future<Output = T::Output> as Pointee>::Metadata {
        let ptr = data.cast::<T>() as NonNull<dyn 'a + Future<Output = T::Output>>;
        ptr.to_raw_parts().1
    }
}

#[repr(C)] // ensure `data` is at the beginning
pub struct Stack<T: ?Sized + Pointee> {
    data: MaybeUninit<[u8; 32]>,
    metadata: T::Metadata,
}
impl<T: ?Sized + Pointee> Stack<T> {
    fn emplace_as<U>(
        constructor: impl FnOnce(VoidPtr) -> NonNull<U>,
        layout: Layout,
    ) -> Option<Self>
    where
        U: ?Sized + Pointee + CoerceInto<T>,
    {
        if layout.align() > std::mem::align_of::<Self>() || layout.size() > 32 {
            return None;
        }

        unsafe {
            let mut data = MaybeUninit::<[u8; 32]>::uninit();
            let data_ptr = NonNull::from(&mut data).cast();
            let output = constructor(data_ptr);
            let metadata = U::coerce(data_ptr, output.to_raw_parts().1);
            Some(Self { data, metadata })
        }
    }
}

// We can now dereference as `&dyn Trait` directly.
impl<T: ?Sized + Pointee> Deref for Stack<T> {
    type Target = T;
    fn deref(&self) -> &Self::Target {
        let ptr = NonNull::<T>::from_raw_parts(NonNull::from(&self.data), self.metadata);
        unsafe { ptr.as_ref() }
    }
}
impl<T: ?Sized + Pointee> DerefMut for Stack<T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        let mut ptr = NonNull::<T>::from_raw_parts(NonNull::from(&mut self.data), self.metadata);
        unsafe { ptr.as_mut() }
    }
}

impl<T: ?Sized + Pointee> Drop for Stack<T> {
    fn drop(&mut self) {
        let ptr = NonNull::<T>::from_raw_parts(NonNull::from(&mut self.data), self.metadata);
        unsafe { ptr.drop_in_place() }
    }
}

// Inherit marker traits.
impl<T: ?Sized + Pointee + Unpin> Unpin for Stack<T> {}
unsafe impl<T: ?Sized + Pointee + Send> Send for Stack<T> {}
unsafe impl<T: ?Sized + Pointee + Sync> Sync for Stack<T> {}

// This implementation may be unnecessary and we can provide `project()` APIs
// instead. Nevertheless, it would be convenient for users.
impl<T: ?Sized + Future> Future for Stack<T> {
    type Output = T::Output;
    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        unsafe { self.map_unchecked_mut(Self::deref_mut).poll(cx) }
    }
}

pub trait Strategy<T: ?Sized + Pointee>: Default {
    type Output;
    fn apply<U>(self, input: impl FnOnce(VoidPtr) -> NonNull<U>, layout: Layout) -> Self::Output
    where
        U: ?Sized + Pointee + CoerceInto<T>;
    //     ^                  ^ This remains here for compatibility with stable dynify.
    //     ` Now we can create the intended container directly from `dyn Trait`
}

// `Strategy` implementation for `Stack`.
#[derive(Default)]
pub struct Stacked;
impl<T: ?Sized + Pointee> Strategy<T> for Stacked {
    type Output = Option<Stack<T>>;
    fn apply<U>(self, input: impl FnOnce(VoidPtr) -> NonNull<U>, layout: Layout) -> Self::Output
    where
        U: ?Sized + Pointee + CoerceInto<T>,
    {
        Stack::emplace_as(input, layout)
    }
}

So far, we have discussed the core design of the new feature, but I haven't verified the practicality of this approach thoroughly, so it may have flaws in some aspects. Additionally, to implement it concretely, several improvements to #[dynify] are still needed:

• Support method-level #[dynify] attributes in a trait item. This adds the flexibility to set a different strategy for a particular method.
• Support in-place transformation of an item. This makes it possible for seamless migration from async_trait to dynify.

Furthermore, we need to add tests as thoroughly as possible since this feature would introduce many unsafe codes.

Footnotes

1. Though I'm opposed to introducing breaking changes since dynify is still in a very early stage. ↩

2. Other solutions are more ugly and would introduce extra runtime overheads. ↩