dirty_fsm/
statemachine.rs

1use std::{collections::HashMap, fmt::Debug, hash::Hash, marker::PhantomData};
2
3use crate::Action;
4use chrono::{DateTime, Utc};
5use log::{debug, trace};
6
7/// Defines a state machine
8#[derive(Debug)]
9pub struct StateMachine<S, Error, Context>
10where
11    S: Debug + Default + Hash + Eq + Clone,
12{
13    current_state: S,
14    last_state: S,
15    last_timestamp: DateTime<Utc>,
16    action_map: HashMap<S, Box<dyn Action<S, Error, Context>>>,
17    _phantom_error: PhantomData<Error>,
18    _phantom_context: PhantomData<Context>,
19}
20
21impl<S, Error, Context> Default for StateMachine<S, Error, Context>
22where
23    S: Debug + Default + Hash + Eq + Clone,
24{
25    fn default() -> Self {
26        Self::new()
27    }
28}
29
30impl<S, Error, Context> StateMachine<S, Error, Context>
31where
32    S: Debug + Default + Hash + Eq + Clone,
33{
34    /// Construct a new `StateMachine`
35    pub fn new() -> Self {
36        Self {
37            current_state: S::default(),
38            last_state: S::default(),
39            last_timestamp: Utc::now(),
40            action_map: HashMap::new(),
41            _phantom_context: PhantomData::default(),
42            _phantom_error: PhantomData::default(),
43        }
44    }
45
46    /// Add an action to the state machine
47    pub fn add_action<A>(&mut self, state: S, mut action: A) -> Result<(), Error>
48    where
49        A: Action<S, Error, Context> + 'static,
50    {
51        debug!("Registering action for state: {:?}", state);
52        #[cfg(feature = "puffin")]
53        puffin::profile_function!();
54
55        // Run the on_register callback
56        action.on_register()?;
57
58        // Actually register the action
59        self.action_map.insert(state, Box::new(action));
60        Ok(())
61    }
62
63    /// Remove an action from the state machine
64    pub fn remove_action(&mut self, state: S) {
65        #[cfg(feature = "puffin")]
66        puffin::profile_function!();
67
68        self.action_map.remove(&state);
69    }
70
71    /// Force-change the current state to something else
72    pub fn force_change_state(&mut self, state: S) -> Result<(), Error> {
73        #[cfg(feature = "puffin")]
74        puffin::profile_function!();
75
76        // End the current state
77        if self.action_map.contains_key(&self.current_state) {
78            // Fetch the current action
79            let action = self.action_map.get_mut(&self.current_state).unwrap();
80
81            // Call the finish function
82            action.on_finish(true)?;
83        }
84
85        debug!("Force-setting state to: {:?}", state);
86        self.current_state = state;
87        Ok(())
88    }
89
90    /// Run a single iteration of the state machine
91    pub fn run(&mut self, context: &Context) -> Result<(), Error> {
92        trace!("Executing a single iteration of the state machine");
93        #[cfg(feature = "puffin")]
94        puffin::profile_function!();
95
96        if self.action_map.contains_key(&self.current_state) {
97            // Fetch the current action
98            let action = self.action_map.get_mut(&self.current_state).unwrap();
99
100            // Handle executing the first run action
101            if self.current_state != self.last_state {
102                debug!(
103                    "Running on_first_run on state {:?} after a state switch",
104                    self.current_state
105                );
106                #[cfg(feature = "puffin")]
107                puffin::profile_scope!("first_run");
108
109                // Execute the first run fn
110                action.on_first_run(context)?;
111            }
112
113            // Calculate the time delta
114            let time_delta = Utc::now() - self.last_timestamp;
115
116            // Execute the action
117            trace!("Executing action for state {:?}", self.current_state);
118            {
119                #[cfg(feature = "puffin")]
120                puffin::profile_scope!("execute");
121
122                let control = action.execute(&time_delta, context)?;
123
124                // Handle the control flags
125                self.last_state = self.current_state.clone();
126                match control {
127                    crate::action::ActionFlag::Continue => {
128                        trace!("Action requested to continue executing");
129                    }
130                    crate::action::ActionFlag::Stop => {
131                        #[cfg(feature = "puffin")]
132                        puffin::profile_scope!("on_stop");
133
134                        trace!("Action requested to stop executing");
135                        self.current_state = S::default();
136                        action.on_finish(false)?;
137                    }
138                    crate::action::ActionFlag::SwitchState(new_state) => {
139                        #[cfg(feature = "puffin")]
140                        puffin::profile_scope!("on_switch");
141
142                        trace!("Action requested to switch to state {:?}", new_state);
143                        self.current_state = new_state;
144                        action.on_finish(false)?;
145                    }
146                }
147            }
148
149            // Update the last timestamp
150            self.last_timestamp = Utc::now();
151        } else {
152            trace!(
153                "No action is configured for state {:?}. Doing nothing",
154                self.current_state
155            );
156        }
157
158        Ok(())
159    }
160}