#![cfg_attr(not(test), no_std)]
#![warn(missing_docs)]
use core::fmt::Debug;
use core::fmt::Display;
use core::future::Future;
use core::marker::PhantomData;
use core::mem;
use core::mem::MaybeUninit;
use core::pin::Pin;
use core::ptr;
use core::task::Context;
use core::task::Poll;
#[cfg(feature = "alloc")]
extern crate alloc;
#[cfg(feature = "alloc")]
use alloc::boxed::Box;
#[repr(C)] pub struct StackFuture<'a, T, const STACK_SIZE: usize> {
data: [MaybeUninit<u8>; STACK_SIZE],
poll_fn: fn(this: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T>,
drop_fn: fn(this: &mut Self),
_phantom: PhantomData<dyn Future<Output = T> + Send + 'a>,
}
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, {
#[allow(clippy::let_unit_value)]
let _ = AssertFits::<F, STACK_SIZE>::ASSERT;
Self::try_from(future).unwrap()
}
pub fn try_from<F>(future: F) -> Result<Self, IntoStackFutureError<F>>
where
F: Future<Output = T> + Send + 'a, {
if Self::has_space_for_val(&future) && Self::has_alignment_for_val(&future) {
let mut result = StackFuture {
data: [MaybeUninit::uninit(); STACK_SIZE],
poll_fn: Self::poll_inner::<F>,
drop_fn: Self::drop_inner::<F>,
_phantom: PhantomData,
};
assert_eq!(result.data.as_ptr() as usize, &result as *const _ as usize);
unsafe { result.as_mut_ptr::<F>().write(future) };
Ok(result)
} else {
Err(IntoStackFutureError::new::<Self>(future))
}
}
#[cfg(feature = "alloc")]
pub fn from_or_box<F>(future: F) -> Self
where
F: Future<Output = T> + Send + 'a, {
Self::try_from(future).unwrap_or_else(|err| Self::from(Box::pin(err.into_inner())))
}
fn poll_inner<F: Future>(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
self.as_pin_mut_ref::<F>().poll(cx)
}
fn drop_inner<F>(&mut self) {
unsafe { ptr::drop_in_place(self.as_mut_ptr::<F>()) }
}
fn as_mut_ptr<F>(&mut self) -> *mut F {
assert!(Self::has_space_for::<F>());
unsafe { mem::transmute(self) }
}
fn as_pin_mut_ref<F>(self: Pin<&mut Self>) -> Pin<&mut F> {
unsafe { self.map_unchecked_mut(|this| &mut *this.as_mut_ptr()) }
}
const fn required_space<F>() -> usize {
mem::size_of::<F>()
}
pub const fn has_space_for<F>() -> bool {
Self::required_space::<F>() <= STACK_SIZE
}
pub const fn has_space_for_val<F>(_: &F) -> bool {
Self::has_space_for::<F>()
}
pub const fn has_alignment_for<F>() -> bool {
mem::align_of::<F>() <= mem::align_of::<Self>()
}
pub const fn has_alignment_for_val<F>(_: &F) -> bool {
Self::has_alignment_for::<F>()
}
}
impl<'a, T, const STACK_SIZE: usize> Future for StackFuture<'a, T, { STACK_SIZE }> {
type Output = T;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
unsafe {
let this = self.get_unchecked_mut();
(this.poll_fn)(Pin::new_unchecked(this), cx)
}
}
}
impl<'a, T, const STACK_SIZE: usize> Drop for StackFuture<'a, T, { STACK_SIZE }> {
fn drop(&mut self) {
(self.drop_fn)(self);
}
}
struct AssertFits<F, const STACK_SIZE: usize>(PhantomData<F>);
impl<F, const STACK_SIZE: usize> AssertFits<F, STACK_SIZE> {
const ASSERT: () = {
if !StackFuture::<F, STACK_SIZE>::has_space_for::<F>() {
panic!("F is too large");
}
if !StackFuture::<F, STACK_SIZE>::has_alignment_for::<F>() {
panic!("F has incompatible alignment");
}
};
}
pub struct IntoStackFutureError<F> {
maximum_size: usize,
maximum_alignment: usize,
future: F,
}
impl<F> IntoStackFutureError<F> {
fn new<Target>(future: F) -> Self {
Self {
maximum_size: mem::size_of::<Target>(),
maximum_alignment: mem::align_of::<Target>(),
future,
}
}
pub fn insufficient_space(&self) -> bool {
self.maximum_size < mem::size_of_val(&self.future)
}
pub fn alignment_too_small(&self) -> bool {
self.maximum_alignment < mem::align_of_val(&self.future)
}
pub fn required_alignment(&self) -> usize {
mem::align_of_val(&self.future)
}
pub fn required_space(&self) -> usize {
mem::size_of_val(&self.future)
}
pub const fn available_alignment(&self) -> usize {
self.maximum_alignment
}
pub const fn available_space(&self) -> usize {
self.maximum_size
}
fn into_inner(self) -> F {
self.future
}
}
impl<F> Display for IntoStackFutureError<F> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
match (self.alignment_too_small(), self.insufficient_space()) {
(true, true) => write!(f,
"cannot create StackFuture, required size is {}, available space is {}; required alignment is {} but maximum alignment is {}",
self.required_space(),
self.available_space(),
self.required_alignment(),
self.available_alignment()
),
(true, false) => write!(f,
"cannot create StackFuture, required alignment is {} but maximum alignment is {}",
self.required_alignment(),
self.available_alignment()
),
(false, true) => write!(f,
"cannot create StackFuture, required size is {}, available space is {}",
self.required_space(),
self.available_space()
),
(false, false) => unreachable!(),
}
}
}
impl<F> Debug for IntoStackFutureError<F> {
fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
f.debug_struct("IntoStackFutureError")
.field("maximum_size", &self.maximum_size)
.field("maximum_alignment", &self.maximum_alignment)
.field("future", &core::any::type_name::<F>())
.finish()
}
}
#[cfg(test)]
mod tests;