vmm/utils/
sm.rs

1// Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
2// SPDX-License-Identifier: Apache-2.0
3
4use std::fmt::Debug;
5
6/// Simple abstraction of a state machine.
7///
8/// `StateMachine<T>` is a wrapper over `T` that also encodes state information for `T`.
9///
10/// Each state for `T` is represented by a `StateFn<T>` which is a function that acts as
11/// the state handler for that particular state of `T`.
12///
13/// `StateFn<T>` returns exactly one other `StateMachine<T>` thus each state gets clearly
14/// defined transitions to other states.
15pub struct StateMachine<T> {
16    function: Option<StateFn<T>>,
17}
18impl<T> Debug for StateMachine<T> {
19    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
20        f.debug_struct("StateMachine")
21            .field("function", &self.function.map(|f| f as usize))
22            .finish()
23    }
24}
25
26/// Type representing a state handler of a `StateMachine<T>` machine. Each state handler
27/// is a function from `T` that handles a specific state of `T`.
28type StateFn<T> = fn(&mut T) -> StateMachine<T>;
29
30impl<T: Debug> StateMachine<T> {
31    /// Creates a new state wrapper.
32    ///
33    /// # Arguments
34    ///
35    /// `function` - the state handler for this state.
36    pub fn new(function: Option<StateFn<T>>) -> StateMachine<T> {
37        StateMachine { function }
38    }
39
40    /// Creates a new state wrapper that has further possible transitions.
41    ///
42    /// # Arguments
43    ///
44    /// `function` - the state handler for this state.
45    pub fn next(function: StateFn<T>) -> StateMachine<T> {
46        StateMachine::new(Some(function))
47    }
48
49    /// Creates a new state wrapper that has no further transitions. The state machine
50    /// will finish after running this handler.
51    ///
52    /// # Arguments
53    ///
54    /// `function` - the state handler for this last state.
55    pub fn finish() -> StateMachine<T> {
56        StateMachine::new(None)
57    }
58
59    /// Runs a state machine for `T` starting from the provided state.
60    ///
61    /// # Arguments
62    ///
63    /// `machine` - a mutable reference to the object running through the various states.
64    /// `starting_state_fn` - a `fn(&mut T) -> StateMachine<T>` that should be the handler for
65    ///                       the initial state.
66    pub fn run(machine: &mut T, starting_state_fn: StateFn<T>) {
67        // Start off in the `starting_state` state.
68        let mut state_machine = StateMachine::new(Some(starting_state_fn));
69        // While current state is not a final/end state, keep churning.
70        while let Some(state_fn) = state_machine.function {
71            // Run the current state handler, and get the next one.
72            state_machine = state_fn(machine);
73        }
74    }
75}
76
77#[cfg(test)]
78mod tests {
79    use super::*;
80
81    // DummyMachine with states `s1`, `s2` and `s3`.
82    #[derive(Debug)]
83    struct DummyMachine {
84        private_data_s1: bool,
85        private_data_s2: bool,
86        private_data_s3: bool,
87    }
88
89    impl DummyMachine {
90        fn new() -> Self {
91            DummyMachine {
92                private_data_s1: false,
93                private_data_s2: false,
94                private_data_s3: false,
95            }
96        }
97
98        // DummyMachine functions here.
99
100        // Simple state-machine: start->s1->s2->s3->done.
101        fn run(&mut self) {
102            // Verify the machine has not run yet.
103            assert!(!self.private_data_s1);
104            assert!(!self.private_data_s2);
105            assert!(!self.private_data_s3);
106
107            // Run the state-machine.
108            StateMachine::run(self, Self::s1);
109
110            // Verify the machine went through all states.
111            assert!(self.private_data_s1);
112            assert!(self.private_data_s2);
113            assert!(self.private_data_s3);
114        }
115
116        fn s1(&mut self) -> StateMachine<Self> {
117            // Verify private data mutates along with the states.
118            assert!(!self.private_data_s1);
119            self.private_data_s1 = true;
120            StateMachine::next(Self::s2)
121        }
122
123        fn s2(&mut self) -> StateMachine<Self> {
124            // Verify private data mutates along with the states.
125            assert!(!self.private_data_s2);
126            self.private_data_s2 = true;
127            StateMachine::next(Self::s3)
128        }
129
130        fn s3(&mut self) -> StateMachine<Self> {
131            // Verify private data mutates along with the states.
132            assert!(!self.private_data_s3);
133            self.private_data_s3 = true;
134            // The machine ends here, adding `s1` as next state to validate this.
135            StateMachine::finish()
136        }
137    }
138
139    #[test]
140    fn test_sm() {
141        let mut machine = DummyMachine::new();
142        machine.run();
143    }
144}