coroutine_state/

lib.rs

//! Inspect the state of a Rust [`Future`] created by an `async` function.
#![no_std]
#![warn(clippy::pedantic)]

use core::{
    future::Future,
    hash::{Hash, Hasher},
    mem::Discriminant,
};

struct DiscriminantExtractor {
    discriminant: Result<Option<u32>, ()>,
}

impl Hasher for DiscriminantExtractor {
    fn finish(&self) -> u64 {
        unreachable!()
    }

    fn write(&mut self, _bytes: &[u8]) {
        self.discriminant = Err(());
    }

    fn write_u32(&mut self, i: u32) {
        if self.discriminant == Ok(None) {
            self.discriminant = Ok(Some(i));
        } else {
            self.discriminant = Err(());
        }
    }
}

fn discriminant_value<T>(discriminant: Discriminant<T>) -> Option<u32> {
    let mut hasher = DiscriminantExtractor {
        discriminant: Ok(None),
    };
    discriminant.hash(&mut hasher);
    hasher.discriminant.ok()?
}

/// State of a coroutine.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum CoroutineState {
    /// The coroutine has not started running.
    Unresumed,
    /// The coroutine has finished running.
    Returned,
    /// The coroutine has panicked.
    Panicked,
    /// The coroutine is suspended at a specific suspension point.
    ///
    /// The suspension points are consecutively numbered starting from 0.
    Suspend(u32),
}

/// Get the current coroutine state of a [`Future`].
///
/// ```
/// use coroutine_state::{coroutine_state, CoroutineState};
///
/// async fn f() {}
///
/// let fut = f();
/// assert_eq!(coroutine_state(&fut), CoroutineState::Unresumed);
/// ```
///
/// # Panics
///
/// You should only pass in a reference to a [`Future`] returned by an `async`
/// function. If you pass in a reference to another type of [`Future`],
/// this function may either panic or return an arbitrary value.
/// However the implemenation uses no `unsafe` code so this can never cause
/// undefined behavior.
///
/// You should be careful not to pass a [`& Pin<&mut Fut>`](core::pin::Pin)
/// instead of a reference to the inner `Fut`:
/// ```
/// use std::pin::pin;
/// use coroutine_state::{coroutine_state, CoroutineState};
///
/// async fn f() {}
///
/// let fut = f();
/// let fut = pin!(fut);
///
/// // This might panic or return an arbitrary state:
/// // println!("{:?}", coroutine_state(&fut));
///
/// // Instead use:
/// println!("{:?}", coroutine_state(&*fut));
/// ```
///
/// Example of where an arbitrary value is returned:
/// ```
/// use std::future::Future;
/// use std::pin::{pin, Pin};
/// use std::task::{Context, Poll};
/// use coroutine_state::{coroutine_state, CoroutineState};
///
/// #[repr(u32)]
/// enum Foo {
///     Foo,
/// }
///
/// impl Future for Foo {
///     type Output = ();
///
///
///     fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
///         Poll::Ready(())
///     }
/// }
///
/// let fut = Foo::Foo;
/// println!("{:?}", coroutine_state(&fut));
/// ```
pub fn coroutine_state<Fut: Future>(fut: &Fut) -> CoroutineState {
    let Some(value) = discriminant_value(core::mem::discriminant(fut)) else {
        panic!(
            "couln't get coroutine state of Future with type {}",
            core::any::type_name::<Fut>()
        );
    };
    match value {
        0 => CoroutineState::Unresumed,
        1 => CoroutineState::Returned,
        2 => CoroutineState::Panicked,
        _ => CoroutineState::Suspend(value - 3),
    }
}

#[cfg(test)]
mod test {
    use core::{
        panic::AssertUnwindSafe,
        pin::{pin, Pin},
        task::{Context, Poll, RawWaker, RawWakerVTable, Waker},
    };

    use super::*;

    extern crate std;
    use std::{prelude::rust_2021::*, vec};

    struct PendingOnce {
        is_ready: bool,
    }

    impl Future for PendingOnce {
        type Output = ();

        fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
            if self.is_ready {
                Poll::Ready(())
            } else {
                self.is_ready = true;
                cx.waker().wake_by_ref();
                Poll::Pending
            }
        }
    }

    fn async_yield() -> impl Future<Output = ()> {
        PendingOnce { is_ready: false }
    }

    const NOOP_RAW_WAKER: RawWaker = {
        const VTABLE: RawWakerVTable =
            RawWakerVTable::new(|_| NOOP_RAW_WAKER, |_| {}, |_| {}, |_| {});
        RawWaker::new(core::ptr::null(), &VTABLE)
    };

    static NOOP_WAKER: Waker = unsafe { Waker::from_raw(NOOP_RAW_WAKER) };
    const NOOP_CONTEXT: Context = Context::from_waker(&NOOP_WAKER);

    fn coroutine_states_inner<Fut: Future>(
        mut fut: Pin<&mut Fut>,
        states: &mut Vec<CoroutineState>,
    ) {
        states.push(coroutine_state(&*fut));

        let mut cx = NOOP_CONTEXT;
        loop {
            let is_ready = fut.as_mut().poll(&mut cx).is_ready();
            states.push(coroutine_state(&*fut));
            if is_ready {
                break;
            }
        }
    }

    fn coroutine_states<Fut: Future>(fut: Fut) -> Vec<CoroutineState> {
        let mut states = vec![];
        let fut = pin!(fut);
        coroutine_states_inner(fut, &mut states);
        states
    }

    #[test]
    fn coroutine_state_empty() {
        #[allow(clippy::unused_async)]
        async fn f() {}

        assert_eq!(
            coroutine_states(f()),
            vec![CoroutineState::Unresumed, CoroutineState::Returned]
        );
    }

    #[test]
    fn coroutine_state_one_yield() {
        async fn f() {
            async_yield().await;
        }

        assert_eq!(
            coroutine_states(f()),
            vec![
                CoroutineState::Unresumed,
                CoroutineState::Suspend(0),
                CoroutineState::Returned,
            ]
        );
    }

    #[test]
    fn coroutine_state_two_yields() {
        async fn f() {
            async_yield().await;
            async_yield().await;
        }

        assert_eq!(
            coroutine_states(f()),
            vec![
                CoroutineState::Unresumed,
                CoroutineState::Suspend(0),
                CoroutineState::Suspend(1),
                CoroutineState::Returned,
            ]
        );
    }

    #[test]
    fn coroutine_state_panic() {
        #[allow(clippy::unused_async)]
        async fn f() {
            panic!();
        }

        let mut states = vec![];
        let fut = f();
        let mut fut = pin!(fut);

        let r = {
            let fut = fut.as_mut();
            std::panic::catch_unwind(AssertUnwindSafe(|| {
                coroutine_states_inner(fut, &mut states);
            }))
        };

        assert!(r.is_err());
        assert_eq!(states, vec![CoroutineState::Unresumed]);

        assert_eq!(coroutine_state(&*fut), CoroutineState::Panicked);
    }
}