vuot 0.0.1

use std::cell::Cell;
use std::future::Future;
use std::marker::PhantomData;
use std::pin::Pin;
use std::ptr::NonNull;
use std::task::{Context, Poll};

use crate::allocator::Alloc;
use crate::frame::{Frame, FrameList};
use crate::sender::Sender;
use crate::StacklessFn;

/// Run the given function with a virtual [`Stack`] that can be used to call
/// functions without growing the actual call stack.
///
/// # Examples
///
/// `run` expects an asynchronous function which takes, as its only parameter,
/// an argument of type [`Stack<'_>`].  This can be done with a normal
/// `async fn`:
///
/// ```
/// # pollster::block_on(async {
/// use vuot::{run, Stack};
/// async fn sum(_: Stack<'_>) -> usize {
///     1 + 2
/// }
/// let result = run(sum).await;
/// assert_eq!(3, result);
/// # })
/// ```
///
/// The [`Stack`] argument represents a handle to a "virtual" call stack.  We
/// use its method [`run`](Stack::run) to poll a future on this virtual call
/// stack.
///
/// ```
/// # pollster::block_on(async {
/// use vuot::{run, Stack};
/// async fn fib(stack: Stack<'_>, n: usize) -> usize {
///     if n < 2 {
///         n
///     } else {
///         stack.run(fib(stack, n - 1)).await + stack.run(fib(stack, n - 2)).await
///     }
/// }
///
/// async fn go(stack: Stack<'_>) -> usize {
///     fib(stack, 10).await
/// }
///
/// let result = run(go).await;
/// assert_eq!(55, result);
/// })
/// ```
///
/// Often, you want to pass some other data to this function.  Due to
/// limitations in Rust, we unfortunately cannot use closures for this:
///
/// ```compile_fail
/// # pollster::block_on(async {
/// # use vuot::{run, Stack};
/// # async fn fib(stack: Stack<'_>, n: usize) -> usize {
/// #     if n < 2 {
/// #         n
/// #     } else {
/// #       stack.run(fib(stack, n - 1)).await + stack.run(fib(stack, n - 2)).await
/// #     }
/// # }
/// let n = 10;
/// let result = run(|stack| fib(stack, n)).await;
/// assert_eq!(55, result);
/// # })
/// ```
///
/// Until Rust gains lending or async closures, you will have to manually create
/// a type to hold any data you want to access and have it implement the
/// [`StacklessFn`] trait.
///
/// ```
/// # pollster::block_on(async {
/// # use std::future::Future;
/// # use vuot::{run, Stack};
/// # async fn fib(stack: Stack<'_>, n: usize) -> usize {
/// #     if n < 2 {
/// #         n
/// #     } else {
/// #       stack.run(fib(stack, n - 1)).await + stack.run(fib(stack, n - 2)).await
/// #     }
/// # }
/// use vuot::StacklessFn;
/// struct Fib<'borrow>(&'borrow usize);
/// impl<'a> StacklessFn<'a, usize> for Fib<'_> {
///     fn call(self, stack: Stack<'a>) -> impl Future<Output = usize> {
///         fib(stack, *self.0)
///     }
/// }
///
/// let n = 10;
/// let result = run(Fib(&n)).await;
/// assert_eq!(55, result);
/// # })
/// ```
pub async fn run<F, T>(f: F) -> T
where
    for<'a> F: StacklessFn<'a, T>,
{
    let alloc = Alloc::new();
    let stack = FrameList::new(&alloc);

    let result = Cell::new(None);

    let stack_ref = Stack { frames: &stack };

    // Note: we don't actually invoke `f.call` until the async block we're
    // setting as the head is being polled.  This is important: it means that
    // `f` will only ever see a stack with itself as the head!
    stack.set_head(async { result.set(Some(f.call(stack_ref).await)) });

    StackPoller { stack: &stack }.await;
    result.into_inner().unwrap()
}

struct StackPoller<'list, 'alloc> {
    stack: &'list FrameList<'alloc>,
}

impl Future for StackPoller<'_, '_> {
    type Output = ();

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // Safety: the only methods on the stack that may be called during this
        // polling are `push` and `consume_frame` from within `Stack`
        unsafe { self.stack.poll(cx) }
    }
}

impl Drop for StackPoller<'_, '_> {
    fn drop(&mut self) {
        // Safety: since the stack poller is being dropped, the entire `run`
        // future is being dropped: nothing else will happen.
        unsafe { self.stack.drop() }
    }
}

/// This is a "handle" to the virtual call stack where futures are stored and
/// polled.
///
/// Note that `Stack` is neither `Send` nor `Sync`, so everything that uses it
/// must be on the same thread as the [`run`] future.
///
/// ```compile_fail,E0277
/// # use vuot::Stack;
/// # async fn go(stack: Stack<'_>) {
/// fn expects_send<T: Send>(t: T) {}
/// expects_send(stack);
/// # }
/// ```
///
/// ```compile_fail,E0277
/// # use vuot::Stack;
/// # async fn go(stack: Stack<'_>) {
/// fn expects_sync<T: Sync>(t: T) {}
/// expects_sync(stack);
/// # }
/// ```
///
/// Still, the future that is returned by [`run`] is `Send`, and so the entire
/// stack can be run in parallell with everything else.
///
/// [moro is not parallel]: https://github.com/nikomatsakis/moro/blob/aafa21566537cd71db666c5e5b27eb69bda6b937/README.md#why-do-moro-spawns-only-run-concurrently-not-parallel
#[derive(Clone, Copy)]
pub struct Stack<'a> {
    frames: &'a FrameList<'a>,
}

impl<'a> Stack<'a> {
    /// Run a future `f` on the virtual call stack.
    ///
    /// Note that the future returned by this function will resolve only once
    /// `f` has returned.  Polling this future before `f` has had a chance to
    /// do so is almost surely a mistake, so this will panic instead.
    ///
    /// If the future `self.run` returns is dropped, the invoked future will
    /// leak.
    ///
    /// As with other functions, you can reference stack-local data in the
    /// future:
    ///
    /// ```
    /// # async fn go(stack: vuot::Stack<'_>) {
    /// let a = 10;
    /// let b = 20;
    /// let (a, b) = stack.run(async { (&a, &b) }).await;
    /// assert_eq!(30, *a + *b);
    /// # }
    /// # pollster::block_on(vuot::run(go));
    /// ```
    ///
    /// And you are not able to leak data out of an inner function:
    ///
    /// ```compile_fail,E0515
    /// # async fn go(stack: vuot::Stack<'_>) {
    /// let (a, b) = stack.run(async {
    ///     let a = 10;
    ///     let b = 20;
    ///     (&a, &b)
    /// }).await;
    /// assert_eq!(30, *a + *b);
    /// # }
    /// # pollster::block_on(vuot::run(go));
    /// ```
    pub fn run<Fut, T>(&self, future: Fut) -> impl Future<Output = T> + 'a
    where
        Fut: Future<Output = T> + 'a,
        T: 'a,
    {
        // Each `run` call should only perform a single heap allocation, that
        // being the allocation of the stack frame. This is a bit tricky, as it
        // turns out:
        //
        // - we need somewhere to store the result `Cell`
        // - we need to construct the heap allocation before returning from this
        //   function
        //
        // Previously, the result was stored in the future returned by this
        // method, with a requirement that the method be polled to completion
        // by the caller to ensure the variable lives long enough.  As it
        // happens, that's not enough: even when using a macro that immediately
        // `.await`ed the future, callers could always just wrap it in another
        // `async { ... }` block to make it suspendable.  UB using safe
        // operations means the unsafe parts are wrong, etc.
        //
        // A better solution is to store the result in the `Sender` itself such
        // that we always have _something_ to write to.  Then, to make sure that
        // is still alive when the returned future is polled, we manage the
        // allocation of the `Sender` here.  This does have the unfortunate side
        // effect that we leak memory if the caller drops our future, which is
        // not great.  I think there is a way to fix that by doing some stack
        // manipulation shenanigans in the `Drop` impl for a custom future this
        // could return, but I'm not sure.

        // Safety: the result pointer is written once the returned future is
        // first polled, which must happen before the caller can yield to `run`
        // (which is when the sender may be polled).
        let sender = unsafe { Sender::new(future) };

        let frames = self.frames;

        // Safety: The returned pointer will not be passed to `consume_frame`
        // until `frame.is_done()`
        let frame = unsafe { frames.push(sender) };

        Receiver {
            frames,
            frame: Cell::new(Some(frame)),
            t: PhantomData,
        }
    }
}

struct Receiver<'a, T> {
    frames: &'a FrameList<'a>,
    frame: Cell<Option<NonNull<Frame>>>,
    t: PhantomData<T>,
}

impl<'a, T> Future for Receiver<'a, T> {
    type Output = T;

    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        // The `FrameList` wakes when we push
        let _ = cx;

        let frame = self.frame.replace(None).expect("polled after completing");

        // Safety: the pointer comes from `stack.push()`, so it is good.
        let frame_ref = unsafe { frame.as_ref() };

        if frame_ref.is_done() {
            // Safety: the pointer comes from `stack.push()` and `is_done()`
            // returned true.
            let frame = unsafe { self.frames.consume_frame(frame) };
            let sender = frame.as_future();

            // Safety: `sender` comes from a `Sender<T, Fut>`, and the
            // pointer is used only until we return, before the box is
            // dropped.
            let result: NonNull<Cell<Option<T>>> =
                unsafe { std::mem::transmute(sender.get_result_pointer()) };

            // Safety: `get_result_pointer` returns a good reference.
            let result = unsafe { result.as_ref() };

            Poll::Ready(result.replace(None).unwrap())
        } else {
            self.frame.set(Some(frame));
            Poll::Pending
        }
    }
}

impl<'a, T> Drop for Receiver<'a, T> {
    fn drop(&mut self) {
        // If the entire call stack is being dropped, then `FrameList` deals
        // with the allocations.
        if self.frames.is_dropping() {
            return;
        }

        // If `self.frame` is `None`, then it returned normally and this future
        // was polled to completion.
        let Some(frame) = self.frame.replace(None) else {
            return;
        };

        // Otherwise, we are being dropped before we have been polled to
        // completion.  There are two situations where this can happen:
        // - the sender was just pushed, but we have not yet yielded to the
        //   stack poller
        // - the sender just completed
        // In the first case, the sender is at the top of the stack, while in
        // the second, we are.  As such, if it is not done, we need to pop it
        // from the stack.

        // Safety: since we're in charge of its allocation, and we have a frame
        // then this is a good pointer.
        let frame_ref = unsafe { frame.as_ref() };
        if !frame_ref.is_done() {
            // Safety: since the frame is not done, it is still at the top of
            // the stack (and nothing else is above it, since the receiver is
            // never polled concurrently with its callees).
            let popped = unsafe { self.frames.pop_head() };
            debug_assert_eq!(frame, popped);
        }

        // Safety: the frame only ever comes from `FrameList::push`
        let frame = unsafe { self.frames.consume_frame(frame) };
        drop(frame);
    }
}