Skip to content

feat: support non-Send stackfutures #34

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
joseph-gio wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: support non-Send stackfutures #34

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
213 changes: 161 additions & 52 deletions src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,26 +75,53 @@ use alloc::boxed::Box;
/// it also respects any alignment requirements for the wrapped future. Note that the
/// wrapped future's alignment must be less than or equal to that of the overall
/// `StackFuture` struct.
// NOTE: we use a type alias rather than a default const generic argument, as that would make methods
// like StackFuture::new ambiguous when calling.
pub type StackFuture<'a, T, const STACK_SIZE: usize> = StackFutureImpl<'a, T, STACK_SIZE, true>;

/// A variant of [`StackFuture`] which allows for futures that do not implement the [`Send`] trait.
///
/// See the documentation of `StackFuture` for more information.
pub type LocalStackFuture<'a, T, const STACK_SIZE: usize> = StackFutureImpl<'a, T, STACK_SIZE, false>;

/// A variant of [`StackFuture`] which supports either [`Send`] ofr non-`Send` futures, depending
/// on the value of the `SEND` const generic argument.
///
/// In most cases, you will want to use `StackFuture` or [`LocalStackFuture`] directly.
/// See the documentation for [`StackFuture`] for more details.
#[repr(C)] // Ensures the data first does not have any padding before it in the struct
pub struct StackFuture<'a, T, const STACK_SIZE: usize> {
pub struct StackFutureImpl<'a, T, const STACK_SIZE: usize, const SEND: bool> {
/// An array of bytes that is used to store the wrapped future.
data: [MaybeUninit<u8>; STACK_SIZE],
/// Since the type of `StackFuture` does not know the underlying future that it is wrapping,
/// we keep a manual vtable that serves pointers to Poll::poll and Drop::drop. These are
/// generated and filled in by `StackFuture::from`.
///
/// This field stores a pointer to the poll function wrapper.
poll_fn: fn(this: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T>,
///
/// SAFETY:
/// * the argument `this` must be the same instance of this type that `poll_fn` was obtained from.
poll_fn: unsafe fn(this: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T>,
/// Stores a pointer to the drop function wrapper
///
/// See the documentation on `poll_fn` for more details.
drop_fn: fn(this: &mut Self),
///
/// SAFETY:
/// * the argument `this` must be the same instance of this type that `drop_fn` was obtained from.
/// * must only be called from the Drop impl of this type.
drop_fn: unsafe fn(this: &mut Self),
/// StackFuture can be used similarly to a `dyn Future`. We keep a PhantomData
/// here so the type system knows this.
_phantom: PhantomData<dyn Future<Output = T> + Send + 'a>,
_phantom: PhantomData<dyn Future<Output = T> + 'a>,
}

impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
// SAFETY:
// We ensure by the API exposed for this type that the contained future will always be Send
// as long as the `SEND` const generic arg is true.
unsafe impl<'a, T, const STACK_SIZE: usize> Send for StackFutureImpl<'a, T, STACK_SIZE, true>
{ }

impl<'a, T, const STACK_SIZE: usize> StackFutureImpl<'a, T, { STACK_SIZE }, true> {
/// Creates a `StackFuture` from an existing future
///
/// See the documentation on [`StackFuture`] for examples of how to use this.
Expand Down Expand Up @@ -135,13 +162,88 @@ impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
/// ```
pub fn from<F>(future: F) -> Self
where
F: Future<Output = T> + Send + 'a, // the bounds here should match those in the _phantom field
F: Future<Output = T> + Send + 'a,
{
Self::from_inner(future)
}

/// Attempts to create a `StackFuture` from an existing future
///
/// If the `StackFuture` is not large enough to hold `future`, this function returns an
/// `Err` with the argument `future` returned to you.
///
/// Panics
///
/// If we cannot satisfy the alignment requirements for `F`, this function will panic.
pub fn try_from<F>(future: F) -> Result<Self, IntoStackFutureError<F>>
where
F: Future<Output = T> + Send + 'a,
{
Self::try_from_inner(future)
}

/// Creates a StackFuture from the given future, boxing if necessary
///
/// This version will succeed even if the future is larger than `STACK_SIZE`. If the future
/// is too large, `from_or_box` will allocate a `Box` on the heap and store the resulting
/// boxed future in the `StackFuture`.
///
/// The same thing also happens if the wrapped future's alignment is larger than StackFuture's
/// alignment.
///
/// This function requires the "alloc" crate feature.
#[cfg(feature = "alloc")]
pub fn from_or_box<F>(future: F) -> Self
where
F: Future<Output = T> + Send + 'a,
{
Self::from_or_box_inner(future)
}
}

impl<'a, T, const STACK_SIZE: usize> StackFutureImpl<'a, T, STACK_SIZE, false> {
/// Creates a `StackFuture` from an existing future.
///
/// See the documentation of [`StackFuture::from`] for more details.
pub fn from<F>(future: F) -> Self
where
F: Future<Output = T> + 'a, // the bounds here should match those in the _phantom field
{
Self::from_inner(future)
}

/// Attempts to create a `StackFuture` from an existing future.
///
/// See the documentation of [`StackFuture::try_from`] for more details.
pub fn try_from<F>(future: F) -> Result<Self, IntoStackFutureError<F>>
where
F: Future<Output = T> + 'a, // the bounds here should match those in the _phantom field
{
Self::try_from_inner(future)
}

/// Creates a StackFuture from the given future, boxing if necessary
///
/// See the documentation of [`StackFuture::from_or_box`] for more details.
#[cfg(feature = "alloc")]
pub fn from_or_box<F>(future: F) -> Self
where
F: Future<Output = T> + 'a, // the bounds here should match those in the _phantom field
{
Self::from_or_box_inner(future)
}
}

impl<'a, T, const STACK_SIZE: usize, const SEND: bool> StackFutureImpl<'a, T, STACK_SIZE, SEND> {
fn from_inner<F>(future: F) -> Self
where
F: Future<Output = T> + 'a, // the bounds here should match those in the _phantom field
{
// Ideally we would provide this as:
//
// impl<'a, F, const STACK_SIZE: usize> From<F> for StackFuture<'a, F::Output, { STACK_SIZE }>
// where
// F: Future + Send + 'a
// F: Future + 'a
//
// However, libcore provides a blanket `impl<T> From<T> for T`, and since `StackFuture: Future`,
// both impls end up being applicable to do `From<StackFuture> for StackFuture`.
Expand All @@ -150,25 +252,26 @@ impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
#[allow(clippy::let_unit_value)]
let _ = AssertFits::<F, STACK_SIZE>::ASSERT;

Self::try_from(future).unwrap()
Self::try_from_inner(future).unwrap()
}

/// Attempts to create a `StackFuture` from an existing future
///
/// If the `StackFuture` is not large enough to hold `future`, this function returns an
/// `Err` with the argument `future` returned to you.
///
/// Panics
///
/// If we cannot satisfy the alignment requirements for `F`, this function will panic.
pub fn try_from<F>(future: F) -> Result<Self, IntoStackFutureError<F>>
fn try_from_inner<F>(future: F) -> Result<Self, IntoStackFutureError<F>>
where
F: Future<Output = T> + Send + 'a, // the bounds here should match those in the _phantom field
F: Future<Output = T> + 'a, // the bounds here should match those in the _phantom field
{
if Self::has_space_for_val(&future) && Self::has_alignment_for_val(&future) {
let mut result = StackFuture {
let mut result = Self {
data: [MaybeUninit::uninit(); STACK_SIZE],
// SAFETY:
// `poll_inner` and `drop_inner` both require `F` to match the future type
// used to construct `self` here. The invariants on the `poll_fn` and `drop_fn`
// fields require that they are only called using the original instance they were
// obtained from, which ensures that `F` will still match at the point
// they are called.
poll_fn: Self::poll_inner::<F>,
// SAFETY:
// the invariants on `drop_fn` transitively uphold the requirement that `drop_inner`
// is only called from the Drop impl of this type.
drop_fn: Self::drop_inner::<F>,
_phantom: PhantomData,
};
Expand All @@ -192,34 +295,35 @@ impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
}
}

/// Creates a StackFuture from the given future, boxing if necessary
///
/// This version will succeed even if the future is larger than `STACK_SIZE`. If the future
/// is too large, `from_or_box` will allocate a `Box` on the heap and store the resulting
/// boxed future in the `StackFuture`.
///
/// The same thing also happens if the wrapped future's alignment is larger than StackFuture's
/// alignment.
///
/// This function requires the "alloc" crate feature.
#[cfg(feature = "alloc")]
pub fn from_or_box<F>(future: F) -> Self
fn from_or_box_inner<F>(future: F) -> Self
where
F: Future<Output = T> + Send + 'a, // the bounds here should match those in the _phantom field
F: Future<Output = T> + 'a, // the bounds here should match those in the _phantom field
{
Self::try_from(future).unwrap_or_else(|err| Self::from(Box::pin(err.into_inner())))
Self::try_from_inner(future).unwrap_or_else(|err| Self::from_inner(Box::pin(err.into_inner())))
}

/// A wrapper around the inner future's poll function, which we store in the poll_fn field
/// of this struct.
fn poll_inner<F: Future>(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
self.as_pin_mut_ref::<F>().poll(cx)
///
/// SAFETY:
/// * the generic argument `F` must be the exact same type originally used to construct
/// this instance via `try_from_inner`.
unsafe fn poll_inner<F: Future>(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
unsafe { self.as_pin_mut_ref::<F>() }.poll(cx)
}

/// A wrapper around the inner future's drop function, which we store in the drop_fn field
/// of this struct.
fn drop_inner<F>(&mut self) {
// SAFETY: *this.as_mut_ptr() was previously written as type F
///
/// SAFETY:
/// * the generic argument `F` must be the exact same type originally used to construct
/// this instance via `try_from_inner`.
/// * must only be called from the drop impl of this type.
unsafe fn drop_inner<F>(&mut self) {
// SAFETY:
// * this.as_mut_ptr() was previously written as type F
// * caller ensures this will only be called from the drop impl of this type.
unsafe { ptr::drop_in_place(self.as_mut_ptr::<F>()) }
}

Expand All @@ -236,8 +340,12 @@ impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
}

/// Returns a pinned mutable reference to a type F stored in self.data
fn as_pin_mut_ref<F>(self: Pin<&mut Self>) -> Pin<&mut F> {
// SAFETY: `StackFuture` is only created by `StackFuture::from`, which
///
/// SAFETY:
/// * the generic argument `F` must be the exact same type originally used to construct
/// this instance via `try_from_inner`.
unsafe fn as_pin_mut_ref<F>(self: Pin<&mut Self>) -> Pin<&mut F> {
// SAFETY: `StackFuture` is only created by `StackFuture::try_from_inner`, which
// writes an `F` to `self.as_mut_ptr(), so it's okay to cast the `*mut F`
// to an `&mut F` with the same lifetime as `self`.
//
Expand Down Expand Up @@ -277,26 +385,27 @@ impl<'a, T, const STACK_SIZE: usize> StackFuture<'a, T, { STACK_SIZE }> {
}
}

impl<'a, T, const STACK_SIZE: usize> Future for StackFuture<'a, T, { STACK_SIZE }> {
impl<'a, T, const STACK_SIZE: usize, const SEND: bool> Future for StackFutureImpl<'a, T, { STACK_SIZE }, SEND> {
type Output = T;

fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
// SAFETY: This is doing pin projection. We unpin self so we can
// access self.poll_fn, and then re-pin self to pass it into poll_in.
// The part of the struct that needs to be pinned is data, since it
// contains a potentially self-referential future object, but since we
// do not touch that while self is unpinned and we do not move self
// while unpinned we are okay.
unsafe {
let this = self.get_unchecked_mut();
(this.poll_fn)(Pin::new_unchecked(this), cx)
}
fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
// SAFETY:
// `self.poll_fn`` is private and never reassigned, so this must be the same instance
// that `poll_fn` was originally obtained from.
unsafe { (self.poll_fn)(self, cx) }
}
}

impl<'a, T, const STACK_SIZE: usize> Drop for StackFuture<'a, T, { STACK_SIZE }> {
impl<'a, T, const STACK_SIZE: usize, const SEND: bool> Drop for StackFutureImpl<'a, T, { STACK_SIZE }, SEND> {
fn drop(&mut self) {
(self.drop_fn)(self);
// SAFETY:
// * `self.drop_fn`` is private and never reassigned, so this must be the same instance
// that `drop_fn` was originally obtained from.
// * we are calling this from the drop impl of this type, which is the only valid place
// to call `drop_fn`.
unsafe {
(self.drop_fn)(self);
}
}
}

Expand Down