krossbar_state_machine/
state.rs

1use std::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use futures::{Future, FutureExt};
7
8/// Machine state
9pub struct State<St, Ret, Err, Fut>
10where
11    Fut: Future<Output = Result<Ret, Err>> + Send,
12{
13    /// Previos state
14    fut: Pin<Box<dyn Future<Output = Result<St, Err>> + Send>>,
15    /// Current state handler
16    func: fn(St) -> Fut,
17    /// State handle future
18    func_fut: Option<Pin<Box<Fut>>>,
19}
20
21impl<St: 'static, Ret: 'static, Err: 'static, Fut> State<St, Ret, Err, Fut>
22where
23    Fut: Future<Output = Result<Ret, Err>> + Send + 'static,
24{
25    pub(crate) fn chain(
26        fut: Pin<Box<dyn Future<Output = Result<St, Err>> + Send>>,
27        func: fn(St) -> Fut,
28    ) -> Self {
29        Self {
30            fut,
31            func,
32            func_fut: None,
33        }
34    }
35
36    /// Add machine state
37    pub fn then<NRet, NFut>(self, func: fn(Ret) -> NFut) -> State<Ret, NRet, Err, NFut>
38    where
39        NRet: 'static,
40        NFut: Future<Output = Result<NRet, Err>> + Send + 'static,
41    {
42        State::chain(Box::pin(self), func)
43    }
44
45    /// Handle final state result
46    pub fn unwrap<NRet>(
47        self,
48        func: fn(Result<Ret, Err>) -> NRet,
49    ) -> impl Future<Output = NRet> + Send
50    where
51        NRet: 'static,
52    {
53        self.map(move |value| func(value))
54    }
55}
56
57impl<St, Ret, Err, Fut> Future for State<St, Ret, Err, Fut>
58where
59    Fut: Future<Output = Result<Ret, Err>> + Send,
60{
61    type Output = Result<Ret, Err>;
62
63    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
64        // Check if we're already made a future from the function
65        let mut func_fut = match self.func_fut.take() {
66            // Already polled previous state, and created a future from the state function
67            // Use it to act as a self future
68            Some(fut) => fut,
69            // Newly created state. Poll previous state first, and create a future from the
70            // inner fuction to poll
71            _ => {
72                let state = match self.fut.as_mut().poll(cx) {
73                    Poll::Ready(value) => match value {
74                        Ok(state) => state,
75                        Err(e) => return Poll::Ready(Err(e)),
76                    },
77                    Poll::Pending => return Poll::Pending,
78                };
79
80                Box::pin((self.func)(state))
81            }
82        };
83
84        match func_fut.as_mut().poll(cx) {
85            Poll::Ready(value) => Poll::Ready(value),
86            Poll::Pending => {
87                // Save function result future to pull later
88                self.func_fut = Some(func_fut);
89                Poll::Pending
90            }
91        }
92    }
93}
94
95impl<St, Ret, Err, Fut> Unpin for State<St, Ret, Err, Fut> where
96    Fut: Future<Output = Result<Ret, Err>> + Send
97{
98}
krossbar_state_machine/state.rs

krossbar_state_machine/
state.rs
