vuot 0.0.1

Run recursive async functions without overflowing the stack
Documentation
use std::future::{poll_fn, Future, Pending};
use std::panic::catch_unwind;
use std::pin::{pin, Pin};
use std::task::{Context, Poll, RawWaker, RawWakerVTable, Waker};

// Note that pollster deals with wakers properly, which is useful to catch any
// accidental deadlocks due to missing wakes.
use pollster::block_on;

use crate::{run, Stack, StacklessFn};

/// References local data in the future pushed to the stack.
#[test]
fn choose_locals() {
    async fn go(stack: Stack<'_>) -> usize {
        let a = 11;
        let b = 12;
        let r = stack.run(async { (&a, &b) }).await;
        *r.0 + *r.1
    }

    let actual = block_on(run(go));
    assert_eq!(23, actual);
}

/// This test copies a `Stack<'_>`, handing it out to two futures and then
/// polling them concurrently (zip-like).
#[test]
fn copy_stack() {
    async fn zonk(stack: Stack<'_>) -> usize {
        stack.run(async { stack.run(async { 501 }).await }).await
    }

    async fn top(stack: Stack<'_>) -> usize {
        let mut a = pin!(zonk(stack));
        let mut b = pin!(zonk(stack));

        let mut a_result = None;
        let mut b_result = None;

        let (a, b) = poll_fn(|cx| {
            let mut done = true;

            if a_result.is_none() {
                match a.as_mut().poll(cx) {
                    Poll::Pending => done = false,
                    Poll::Ready(value) => a_result = Some(value),
                }
            }

            if b_result.is_none() {
                match b.as_mut().poll(cx) {
                    Poll::Pending => done = false,
                    Poll::Ready(value) => b_result = Some(value),
                }
            }

            if done {
                Poll::Ready((a_result.unwrap(), b_result.unwrap()))
            } else {
                Poll::Pending
            }
        })
        .await;

        a + b
    }

    let actual = block_on(run(top));
    assert_eq!(1002, actual);
}

/// Apparently, Miri has some issues with `bumpalo`, `allocator_api`, and
/// pointer roundtrips (see https://github.com/fitzgen/bumpalo/issues/247).
/// Until that issue is resolved, this test remains disabled when using
/// `bumpalo` and Miri.
#[cfg(not(all(miri, feature = "bumpalo")))]
#[test]
fn double_stack() {
    async fn make1(stack: Stack<'_>) -> usize {
        stack.run(async { 11 }).await
    }

    async fn make(stack: Stack<'_>) -> usize {
        stack.run(make1(stack)).await + stack.run(make1(stack)).await
    }

    async fn go(stack: Stack<'_>) -> usize {
        make(stack).await
    }

    let actual = block_on(run(go));
    assert_eq!(22, actual);
}

/// This one allocates memory a little bit differently from the variant above,
/// so it works fine with `bumpalo` and Miri.
#[test]
fn double_stack_variant() {
    async fn make1(stack: Stack<'_>) -> usize {
        stack.run(async { 54 }).await
    }

    async fn make(stack: Stack<'_>) -> usize {
        let a = stack.run(make1(stack));
        let b = stack.run(make1(stack));
        a.await + b.await
    }

    async fn go(stack: Stack<'_>) -> usize {
        make(stack).await
    }

    let actual = block_on(run(go));
    assert_eq!(108, actual);
}

/// Pushes a future to the stack, yields so it can run to completion, then drops
/// the receiver.
#[test]
fn drop_after_completed() {
    async fn go(stack: Stack<'_>) -> usize {
        {
            let inner = stack.run(async { 853 });
            let mut yielded = false;
            poll_fn(|_| {
                if yielded {
                    Poll::Ready(())
                } else {
                    yielded = true;
                    Poll::Pending
                }
            })
            .await;

            drop(inner);
        }

        198
    }

    let actual = block_on(run(go));
    assert_eq!(198, actual);
}

/// Panics before returning any futures.
#[test]
fn eager_panic() {
    fn go(_: Stack<'_>) -> Pending<()> {
        panic!("all is well");
    }

    let crash = catch_unwind(|| block_on(run(go)));
    match crash {
        Ok(()) => panic!("should panic"),
        Err(e) => {
            assert_eq!(Some(&"all is well"), e.downcast_ref());
        }
    }
}

/// Uses `stack` as soon as we obtain one.  Ensures that user code is never
/// given access to an in-construction `Stack`.
#[test]
fn immediately_pushes() {
    fn go(stack: Stack<'_>) -> impl Future<Output = usize> + '_ {
        let future = stack.run(async { 9821 });
        async move {
            future.await;
            1357
        }
    }

    let actual = block_on(run(go));
    assert_eq!(1357, actual);
}

/// Panics inside some nested futures.
#[test]
fn lazy_panic() {
    async fn go(stack: Stack<'_>) {
        stack
            .run(async {
                stack
                    .run(async { stack.run(async { panic!("ouch!") }).await })
                    .await
            })
            .await
    }

    let crash = catch_unwind(|| block_on(run(go)));
    match crash {
        Ok(()) => panic!("should panic"),
        Err(e) => {
            assert_eq!(Some(&"ouch!"), e.downcast_ref());
        }
    }
}

/// This test recursively traverses and sums up the values of a linked list.
#[test]
fn linked_list() {
    struct List(usize, Option<Box<List>>);

    async fn sum(stack: Stack<'_>, list: &List) -> usize {
        list.0
            + if let Some(next) = &list.1 {
                stack.run(sum(stack, next)).await
            } else {
                0
            }
    }

    struct Fn<'a>(&'a List);
    impl<'a> StacklessFn<'a, usize> for Fn<'_> {
        fn call(self, stack: Stack<'a>) -> impl Future<Output = usize> {
            sum(stack, self.0)
        }
    }

    let mut list = List(0, None);
    for i in 1..=100 {
        list = List(i, Some(Box::new(list)));
    }

    let actual = block_on(run(Fn(&list)));
    assert_eq!(5050, actual);
}

/// This test explicitly nests a bunch of `run` calls.  Unless something has
/// gone wrong, this should produce a value almost immediately.
#[test]
fn nested() {
    async fn go(stack: Stack<'_>) -> usize {
        stack
            .run(async {
                stack
                    .run(async {
                        stack
                            .run(async { stack.run(async { 192837465 }).await })
                            .await
                    })
                    .await
            })
            .await
    }

    let actual = block_on(run(go));
    assert_eq!(192837465, actual);
}

/// Run a future on the stack and keep polling it, even after it has finished.
/// Mostly useful with Miri.
#[test]
#[should_panic]
fn poll_after_done() {
    async fn go(stack: Stack<'_>) -> ! {
        let mut future = pin!(stack.run(async { String::from("hii") }));
        poll_fn(|cx| {
            let _ = future.as_mut().poll(cx);
            cx.waker().wake_by_ref();
            Poll::Pending
        })
        .await
    }

    block_on(run(go));
}

/// A future is pushed, then its receiver is polled once, then the receiver is
/// dropped. Finally, we yield once.
#[test]
fn sneaky_drop() {
    async fn go(stack: Stack<'_>) -> usize {
        {
            let f = pin!(stack.run(async { 101 }));

            let waker = noop();
            let mut context = Context::from_waker(&waker);

            assert_eq!(Poll::Pending, f.poll(&mut context));
        }

        let mut yielded = false;
        poll_fn(|cx| {
            if yielded {
                Poll::Ready(123)
            } else {
                cx.waker().wake_by_ref();
                yielded = true;
                Poll::Pending
            }
        })
        .await
    }

    let actual = block_on(run(go));
    assert_eq!(123, actual);
}

/// This test ensures that stack frames are dropped from top to bottom.  We have
/// an outer future which allocates a value and runs a future `Innermost` that
/// uses the reference in its drop.  This is only sound if the innermost futures
/// have the shortest lifetimes.
#[test]
fn wacky_drop() {
    struct Innermost<'a, 'c>(&'a usize, &'c mut usize);

    impl Future for Innermost<'_, '_> {
        type Output = ();
        fn poll(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll<()> {
            Poll::Pending
        }
    }

    impl Drop for Innermost<'_, '_> {
        fn drop(&mut self) {
            *self.1 = *self.0;
        }
    }

    struct Fn<'c>(&'c mut usize);
    impl<'a> StacklessFn<'a, ()> for Fn<'_> {
        async fn call(self, stack: Stack<'a>) {
            let value = 0xcafebeef;
            stack.run(Innermost(&value, self.0)).await
        }
    }

    let mut result = 0;

    {
        let mut future = pin!(run(Fn(&mut result)));

        let waker = noop();
        let mut context = Context::from_waker(&waker);

        // Poll once to spawn the inner future
        assert_eq!(Poll::Pending, future.as_mut().poll(&mut context));
    }

    assert_eq!(0xcafebeef, result);
}

/// Polls a non-stack future inside one on the stack.
#[test]
fn wakeup() {
    fn yield_once() -> impl Future<Output = ()> {
        let mut yielded = false;
        poll_fn(move |cx| {
            if yielded {
                Poll::Ready(())
            } else {
                yielded = true;
                cx.waker().wake_by_ref();
                Poll::Pending
            }
        })
    }

    async fn go(stack: Stack<'_>) -> &'static str {
        stack
            .run(async {
                yield_once().await;
                "hello after all"
            })
            .await
    }

    let actual = block_on(run(go));
    assert_eq!("hello after all", actual);
}

/// Create a waker that does absolutely nothing.
fn noop() -> Waker {
    const RAW: RawWaker = RawWaker::new(
        core::ptr::null(),
        &RawWakerVTable::new(clone, wake, wake_by_ref, drop),
    );

    unsafe fn clone(_: *const ()) -> RawWaker {
        RAW
    }

    unsafe fn wake(_: *const ()) {}
    unsafe fn wake_by_ref(_: *const ()) {}
    unsafe fn drop(_: *const ()) {}

    // Safety: none of the functions do anything, and `clone` returns the
    // correct raw waker.
    unsafe { Waker::from_raw(RAW) }
}