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}