cps_st_machine/
lib.rs

1//! This crate defines a set of primitives useful in creating state machines
2//! driven by continuation passing style.
3
4mod abort;
5
6use std::hint::unreachable_unchecked;
7use std::panic;
8use std::sync::Once;
9
10use self::abort::{
11    abort_execution, abort_execution_panic_middleware, extract_return_value,
12    is_abort_execution_exception,
13};
14
15/// Executable [`State`].
16pub type ExecutableState<Ctx, Ret> = fn(&mut Ctx) -> State<Ctx, Ret>;
17
18/// Some state in a state machine.
19#[repr(transparent)]
20pub struct State<Ctx, Ret> {
21    inner: for<'ctx> fn(&'ctx mut Ctx) -> Ret,
22}
23
24impl<Ctx, Ret> State<Ctx, Ret> {
25    fn executable(&self) -> ExecutableState<Ctx, Ret> {
26        unsafe { std::mem::transmute(self.inner) }
27    }
28
29    /// Perform a state transition.
30    pub fn next(self, context: &mut Ctx) -> Self {
31        self.executable()(context)
32    }
33
34    /// Bail out of a state machine execution environment.
35    pub fn bail(return_value: Ret) -> Self
36    where
37        Ret: Send + 'static,
38    {
39        abort_execution(return_value)
40    }
41
42    /// Build a new [`State`], by supplying its code.
43    pub fn from_fn(state: ExecutableState<Ctx, Ret>) -> Self {
44        Self {
45            inner: unsafe {
46                std::mem::transmute::<
47                    for<'ctx> fn(&'ctx mut Ctx) -> State<Ctx, Ret>,
48                    for<'ctx> fn(&'ctx mut Ctx) -> Ret,
49                >(state)
50            },
51        }
52    }
53}
54
55/// The entry point to a state machine, whose initial [`State`] is
56/// `init_state`.
57///
58/// # Implementation details
59///
60/// This function defines a special execution environment that
61/// overrides the panic handler. Library users should be aware
62/// of this fact, if they intend to replace the default panic
63/// handler.
64pub fn run<Ctx, Ret>(context: Ctx, init_state: State<Ctx, Ret>) -> Ret
65where
66    Ret: 'static,
67    Ctx: panic::UnwindSafe,
68{
69    static INSTALL_HANDLER: Once = Once::new();
70
71    INSTALL_HANDLER.call_once(|| {
72        let previous_hook = panic::take_hook();
73        panic::set_hook(abort_execution_panic_middleware::<Ret>(previous_hook));
74    });
75
76    // run the state machine until a panic is encountered
77    let result = panic::catch_unwind(|| run_inner(context, init_state));
78
79    // handle the panic
80    match result {
81        // an actual panic
82        Err(exception) if !is_abort_execution_exception::<Ret>(&*exception) => {
83            panic::resume_unwind(exception);
84        }
85        // a panic initiated by the state machine, in order
86        // to abort execution
87        Err(mut abort) => {
88            // SAFETY: we only call `extract_return_value` once, upon returning from
89            // a state machine execution env, and we have asserted this is an abort
90            unsafe { extract_return_value(&mut *abort) }
91        }
92        // SAFETY: `run_inner` never returns
93        Ok(_) => unsafe { unreachable_unchecked() },
94    }
95}
96
97#[inline(always)]
98fn run_inner<Ctx, Ret>(mut context: Ctx, mut state: State<Ctx, Ret>) -> ! {
99    loop {
100        state = state.next(&mut context);
101    }
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn terminate() {
110        assert_eq!(run((), State::from_fn(basic_ret_i32::entry::<1234>)), 1234);
111        assert_eq!(run((), State::from_fn(basic_ret_i32::entry::<42>)), 42);
112    }
113
114    mod basic_ret_i32 {
115        use super::*;
116
117        pub fn entry<const RET: i32>(_: &mut ()) -> State<(), i32> {
118            State::from_fn(middle_state::<RET>)
119        }
120
121        fn middle_state<const RET: i32>(_: &mut ()) -> State<(), i32> {
122            State::from_fn(end_state::<RET>)
123        }
124
125        fn end_state<const RET: i32>(_: &mut ()) -> State<(), i32> {
126            State::bail(RET)
127        }
128    }
129}