ispawn 0.0.2

//! A low-cost, `no_std` abstraction over `Future` spawners that erases the underlying spawn handle's
//! type. This allows writing libraries that spawn futures in an executor-agnostic way without
//! using viral type parameters throughout.
//!
//! ispawn's traits do not require an additional heap allocation per spawn (on top of the heap
//! allocation the underlying executor likely does). While spawning is typically a relatively
//! infrequent operation, and thus not particularly performance-sensitive, boxing a future before
//! spawning adds an extra pointer indirection and virtual dispatch every time the future is polled.
//!
//! Supporting executors expose a way to split spawning into two operations:
//! 1. allocate memory for task
//! 2. write future and queue to executor
//!
//! The first operation allocates without knowing the concrete type of the future - it's only given
//! the layout. Executors tend to wrap spawned futures in a task datastructure, with the future
//! potentially inline as a DST - the first step allocates this task datastructure but leaves the
//! future uninitialized.
//!
//! The second operation writes the future to the task structure and queues it toward the executor.
//!
//! This split allows the `spawn` call to be inlined such that the future's state can be written
//! directly to its final destination - hopefully avoiding a potentially large memcpy from the stack
//! to the heap.
//!
//! # Adding support for an executor
//!
//!

#![no_std]

#[cfg(feature = "alloc")]
extern crate alloc;

#[cfg(feature = "alloc")]
use alloc::rc::Rc;
use core::{
    alloc::Layout,
    future::Future,
};

struct LocalSpawnerVtable {
    spawn_dyn: fn(handle: *const (), future_layout: Layout) -> SpawnCompleter,
    finish_spawn: fn(handle: *const (), task_ptr_as_dyn_future: *mut dyn Future<Output = ()>);
    on_clone: fn(handle: *const ()),
    on_drop: fn(handle: *const ()),
}

impl LocalSpawnerVtable {
    fn get<T: IntoLocalSpawner>() -> &'static Self {
        &LocalSpawnerVtable {
            spawn_dyn: T::spawn_dyn,
            spawn_dyn: T::spawn_dyn,
            on_clone: T::on_clone,
            on_drop: T::on_drop,
        }
    }
}

pub struct LocalSpawner {
    handle: *const (),
    vtable: &'static LocalSpawnerVtable,
}

impl LocalSpawner {
    #[cfg(feature = "alloc")]
    pub fn new<T: IntoLocalSpawner>(inner: T) -> Self {
        /*const VTABLE: LocalSpawnerVtable = LocalSpawnerVtable {
            spawn_fn: T::spawn_dyn,
            on_clone: T::on_clone,
            on_drop: T::on_drop,
        };*/

        Self {
            handle: T::into_raw(inner),
            vtable: LocalSpawnerVtable::get::<T>(),
        }
    }

    pub fn spawn<F: Future<Output = ()> + 'static>(&self, f: F) {
        let spawn_completer = (self.vtable.spawn_dyn)(self.handle, Layout::new::<F>());
        unsafe {
            spawn_completer.spawn(f);
        }
    }
}

impl Clone for LocalSpawner {
    fn clone(&self) -> Self {
        (self.vtable.on_clone)(self.handle);
        Self {
            handle: self.handle,
            vtable: self.vtable,
        }
    }
}

impl Drop for LocalSpawner {
    fn drop(&mut self) {
        (self.vtable.on_drop)(self.handle);
    }
}

/// This trait is intended to be implemented for reference counted spawn handles. If your
/// `T: LocalSpawnDyn` is not `Clone`, it can be wrapped in `Rc` - `Rc<T>` will implement
/// `IntoLocalSpawner` via ispawn's blanket impl.
///
/// Safety: the pointer returned by `Self::into_raw` must be 'static.
pub unsafe trait IntoLocalSpawner {
    fn into_raw(self) -> *const ();
    fn spawn_dyn(handle: *const (), future_layout: Layout) -> SpawnCompleter;
    fn on_clone(handle: *const ());
    fn on_drop(handle: *const ());
}

#[cfg(feature = "alloc")]
unsafe impl<T: LocalSpawnDyn> IntoLocalSpawner for Rc<T> {
    fn into_raw(self) -> *const () {
        Rc::into_raw(self) as *const ()
    }

    fn spawn_dyn<'a>(handle: *const (), future_layout: Layout) -> SpawnCompleter {
        let spawner: &T = &*(handle as *const T);
        LocalSpawnDyn::spawn_dyn(spawner, future_layout)
    }

    fn on_clone(handle: *const ()) {
        Rc::increment_strong_count(handle as *const T);
    }

    fn on_drop(handle: *const ()) {
        drop(Rc::from_raw(handle as *const T));
    }
}

pub trait LocalSpawnDyn {
    /// The `spawn` method must be called on the returned SpawnCompleter. If the `spawn_dyn`
    /// implementation allocates memory for the task (usually the case), not doing so will result
    /// in a memory leak.
    fn spawn_dyn(&self, future_layout: Layout) -> SpawnCompleter;

    unsafe fn finish_spawn(&self, task_ptr_as_dyn_future: *mut dyn Future<Output = ()>);
}

pub struct SpawnCompleter {
    handle: *const (),
    vtable: &'static LocalSpawnerVtable,
    task_ptr: *mut (),
    future_ptr: *mut (),
}

impl SpawnCompleter {
    pub fn new(spawner: &LocalSpawner, task_ptr: *mut (), future_ptr: *mut ()) -> Self {
        Self {
            handle: spawner.handle,
            vtable: spawner.vtable,
            task_ptr,
            future_ptr,
        }
    }

    /// Safety: The caller must ensure that `F` has the same layout that was used to create this
    /// `SpawnCompleter`.
    unsafe fn spawn<F: Future<Output = ()> + 'static>(self, f: F) {
        core::ptr::write(self.future_ptr as *mut F, f);

        // Learned this trick from here:
        //   https://www.reddit.com/r/rust/comments/hcofkh/comment/fvgpv5e
        // This seems pretty dubious, but it works today.
        // TODO elaborate how this works
        self.spawner.finish_spawn(self.task_ptr as *mut F as *mut dyn Future<Output = ()>);
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn it_works() {
        let result = add(2, 2);
        assert_eq!(result, 4);
    }
}