rust_robotics_control/
state_machine.rs

1/*!
2 * State Machine implementation for Mission Planning
3 *
4 * This module implements a finite state machine for robot behavior management.
5 * It provides a flexible framework for defining states, transitions, events,
6 * guards, and actions in robotics applications.
7 *
8 * Ported from PythonRobotics
9 * Original author: Wang Zheng (@Aglargil)
10 */
11
12use std::collections::HashMap;
13use std::fmt;
14
15/// Type alias for callback functions
16pub type CallbackFn = Box<dyn Fn()>;
17pub type GuardFn = Box<dyn Fn() -> bool>;
18pub type ActionFn = Box<dyn Fn()>;
19
20/// Represents a state in the state machine
21#[derive(Clone, PartialEq)]
22pub struct State {
23    pub name: String,
24    on_enter: Option<String>, // Store function name as string for demo purposes
25    on_exit: Option<String>,
26}
27
28impl State {
29    /// Create a new state
30    pub fn new(name: &str) -> Self {
31        State {
32            name: name.to_string(),
33            on_enter: None,
34            on_exit: None,
35        }
36    }
37
38    /// Create a new state with callbacks
39    pub fn with_callbacks(name: &str, on_enter: Option<&str>, on_exit: Option<&str>) -> Self {
40        State {
41            name: name.to_string(),
42            on_enter: on_enter.map(|s| s.to_string()),
43            on_exit: on_exit.map(|s| s.to_string()),
44        }
45    }
46
47    /// Enter the state
48    pub fn enter(&self) {
49        println!("entering <{}>", self.name);
50        if let Some(ref callback) = self.on_enter {
51            println!("  executing on_enter: {}", callback);
52        }
53    }
54
55    /// Exit the state
56    pub fn exit(&self) {
57        println!("exiting <{}>", self.name);
58        if let Some(ref callback) = self.on_exit {
59            println!("  executing on_exit: {}", callback);
60        }
61    }
62}
63
64impl fmt::Display for State {
65    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
66        write!(f, "{}", self.name)
67    }
68}
69
70/// Represents a transition in the state machine
71#[derive(Clone)]
72pub struct Transition {
73    pub src_state: String,
74    pub event: String,
75    pub dst_state: String,
76    pub guard: Option<String>,
77    pub action: Option<String>,
78}
79
80impl Transition {
81    pub fn new(src_state: &str, event: &str, dst_state: &str) -> Self {
82        Transition {
83            src_state: src_state.to_string(),
84            event: event.to_string(),
85            dst_state: dst_state.to_string(),
86            guard: None,
87            action: None,
88        }
89    }
90
91    pub fn with_guard(mut self, guard: &str) -> Self {
92        self.guard = Some(guard.to_string());
93        self
94    }
95
96    pub fn with_action(mut self, action: &str) -> Self {
97        self.action = Some(action.to_string());
98        self
99    }
100}
101
102/// Main State Machine implementation
103pub struct StateMachine {
104    name: String,
105    states: HashMap<String, State>,
106    events: HashMap<String, String>,
107    transitions: HashMap<(String, String), Transition>,
108    current_state: Option<State>,
109    transition_history: Vec<(String, String, String)>, // (from, event, to)
110}
111
112impl StateMachine {
113    /// Create a new state machine
114    pub fn new(name: &str) -> Self {
115        StateMachine {
116            name: name.to_string(),
117            states: HashMap::new(),
118            events: HashMap::new(),
119            transitions: HashMap::new(),
120            current_state: None,
121            transition_history: Vec::new(),
122        }
123    }
124
125    /// Register a state in the state machine
126    pub fn register_state(&mut self, state: State) {
127        self.states.insert(state.name.clone(), state);
128    }
129
130    /// Register an event in the state machine
131    pub fn register_event(&mut self, event: &str) {
132        self.events.insert(event.to_string(), event.to_string());
133    }
134
135    /// Add a transition to the state machine
136    pub fn add_transition(&mut self, transition: Transition) {
137        // Register states and events if they don't exist
138        if !self.states.contains_key(&transition.src_state) {
139            self.register_state(State::new(&transition.src_state));
140        }
141        if !self.states.contains_key(&transition.dst_state) {
142            self.register_state(State::new(&transition.dst_state));
143        }
144        self.register_event(&transition.event);
145
146        let key = (transition.src_state.clone(), transition.event.clone());
147        self.transitions.insert(key, transition);
148    }
149
150    /// Set the initial state
151    pub fn set_initial_state(&mut self, state_name: &str) {
152        if let Some(state) = self.states.get(state_name) {
153            self.current_state = Some(state.clone());
154            println!("|{}| initial state set to <{}>", self.name, state_name);
155            state.enter();
156        } else {
157            panic!("State '{}' not found", state_name);
158        }
159    }
160
161    /// Get the current state
162    pub fn get_current_state(&self) -> Option<&State> {
163        self.current_state.as_ref()
164    }
165
166    /// Process an event
167    pub fn process(&mut self, event: &str) -> Result<(), String> {
168        if let Some(ref current_state) = self.current_state {
169            let key = (current_state.name.clone(), event.to_string());
170
171            if let Some(transition) = self.transitions.get(&key).cloned() {
172                self.execute_transition(&transition, event)
173            } else {
174                Err(format!(
175                    "|{}| invalid transition: <{}> : [{}]",
176                    self.name, current_state.name, event
177                ))
178            }
179        } else {
180            Err("State machine is not initialized".to_string())
181        }
182    }
183
184    /// Execute a transition
185    fn execute_transition(&mut self, transition: &Transition, event: &str) -> Result<(), String> {
186        // Safety: execute_transition is only called from process() which checks current_state is Some
187        let current_state = self
188            .current_state
189            .as_ref()
190            .expect("execute_transition called without current_state");
191
192        // Check guard condition
193        if let Some(ref guard) = transition.guard {
194            println!("  checking guard: {}", guard);
195            // For demo purposes, we'll simulate guard evaluation
196            if !self.evaluate_guard(guard) {
197                println!(
198                    "|{}| skipping transition from <{}> to <{}> because guard [{}] failed",
199                    self.name, current_state.name, transition.dst_state, guard
200                );
201                return Ok(());
202            }
203        }
204
205        // Execute action
206        if let Some(ref action) = transition.action {
207            println!("  executing action: {}", action);
208            self.execute_action(action);
209        }
210
211        // Perform state transition
212        if current_state.name != transition.dst_state {
213            println!(
214                "|{}| transitioning from <{}> to <{}> on event [{}]",
215                self.name, current_state.name, transition.dst_state, event
216            );
217
218            // Record transition history
219            self.transition_history.push((
220                current_state.name.clone(),
221                event.to_string(),
222                transition.dst_state.clone(),
223            ));
224
225            // Exit current state
226            current_state.exit();
227
228            // Enter new state
229            if let Some(new_state) = self.states.get(&transition.dst_state) {
230                self.current_state = Some(new_state.clone());
231                new_state.enter();
232            } else {
233                return Err(format!(
234                    "Destination state '{}' not found",
235                    transition.dst_state
236                ));
237            }
238        } else {
239            println!(
240                "|{}| self-transition on <{}> with event [{}]",
241                self.name, current_state.name, event
242            );
243        }
244
245        Ok(())
246    }
247
248    /// Evaluate a guard condition (simplified for demo)
249    fn evaluate_guard(&self, guard: &str) -> bool {
250        match guard {
251            "can_start" => true,
252            "has_battery" => true,
253            "obstacle_detected" => false,
254            "goal_reached" => false,
255            "emergency" => false,
256            _ => true, // Default to true for unknown guards
257        }
258    }
259
260    /// Execute an action (simplified for demo)
261    fn execute_action(&self, action: &str) {
262        match action {
263            "start_motors" => println!("    -> Starting motors"),
264            "stop_motors" => println!("    -> Stopping motors"),
265            "play_sound" => println!("    -> Playing notification sound"),
266            "save_position" => println!("    -> Saving current position"),
267            "send_alert" => println!("    -> Sending emergency alert"),
268            _ => println!("    -> Executing action: {}", action),
269        }
270    }
271
272    /// Get transition history
273    pub fn get_transition_history(&self) -> &Vec<(String, String, String)> {
274        &self.transition_history
275    }
276
277    /// Generate a simple text representation of the state machine
278    pub fn generate_diagram(&self) -> String {
279        let mut diagram = Vec::new();
280        diagram.push(format!("State Machine: {}", self.name));
281        diagram.push("".to_string());
282
283        if let Some(ref current) = self.current_state {
284            diagram.push(format!("Current State: {}", current.name));
285            diagram.push("".to_string());
286        }
287
288        diagram.push("States:".to_string());
289        for state in self.states.values() {
290            let marker = if Some(state) == self.current_state.as_ref() {
291                " [CURRENT]"
292            } else {
293                ""
294            };
295            diagram.push(format!("  - {}{}", state.name, marker));
296        }
297        diagram.push("".to_string());
298
299        diagram.push("Transitions:".to_string());
300        for transition in self.transitions.values() {
301            let mut trans_str = format!(
302                "  {} --[{}]--> {}",
303                transition.src_state, transition.event, transition.dst_state
304            );
305
306            if let Some(ref guard) = transition.guard {
307                trans_str.push_str(&format!(" [guard: {}]", guard));
308            }
309            if let Some(ref action) = transition.action {
310                trans_str.push_str(&format!(" / {}", action));
311            }
312            diagram.push(trans_str);
313        }
314
315        diagram.join("\n")
316    }
317
318    /// Create a visualization of the state machine as SVG
319    pub fn visualize(&self, filename: &str) {
320        use std::io::Write;
321
322        // Get sorted state names for consistent layout
323        let mut state_names: Vec<_> = self.states.keys().cloned().collect();
324        state_names.sort();
325        let n_states = state_names.len();
326
327        if n_states == 0 {
328            return;
329        }
330
331        // SVG dimensions
332        let width = 640.0;
333        let height = 640.0;
334        let cx = width / 2.0;
335        let cy = height / 2.0;
336        let radius = 200.0; // Circle radius for state layout
337        let node_radius = 40.0;
338
339        // Position states in a circle
340        let mut state_positions: HashMap<String, (f64, f64)> = HashMap::new();
341        for (i, state_name) in state_names.iter().enumerate() {
342            let angle = std::f64::consts::PI / 2.0
343                - 2.0 * std::f64::consts::PI * i as f64 / n_states as f64;
344            let x = cx + radius * angle.cos();
345            let y = cy - radius * angle.sin(); // SVG y is inverted
346            state_positions.insert(state_name.clone(), (x, y));
347        }
348
349        let mut svg = String::new();
350
351        // SVG header
352        svg.push_str(&format!(
353            r##"<?xml version="1.0" encoding="UTF-8"?>
354<svg xmlns="http://www.w3.org/2000/svg" width="{}" height="{}" viewBox="0 0 {} {}">
355  <defs>
356    <marker id="arrowhead" markerWidth="10" markerHeight="7" refX="9" refY="3.5" orient="auto">
357      <polygon points="0 0, 10 3.5, 0 7" fill="#666"/>
358    </marker>
359  </defs>
360  <rect width="100%" height="100%" fill="white"/>
361  <text x="{}" y="30" text-anchor="middle" font-family="Arial" font-size="18" font-weight="bold">Robot State Machine</text>
362"##,
363            width, height, width, height, cx
364        ));
365
366        // Draw transitions as arrows
367        for transition in self.transitions.values() {
368            if let (Some(&(x1, y1)), Some(&(x2, y2))) = (
369                state_positions.get(&transition.src_state),
370                state_positions.get(&transition.dst_state),
371            ) {
372                if transition.src_state != transition.dst_state {
373                    let dx = x2 - x1;
374                    let dy = y2 - y1;
375                    let dist = (dx * dx + dy * dy).sqrt();
376
377                    // Offset from node centers
378                    let start_x = x1 + (node_radius + 5.0) * dx / dist;
379                    let start_y = y1 + (node_radius + 5.0) * dy / dist;
380                    let end_x = x2 - (node_radius + 10.0) * dx / dist;
381                    let end_y = y2 - (node_radius + 10.0) * dy / dist;
382
383                    // Draw arrow line
384                    svg.push_str(&format!(
385                        r##"  <line x1="{:.1}" y1="{:.1}" x2="{:.1}" y2="{:.1}" stroke="#666" stroke-width="2" marker-end="url(#arrowhead)"/>
386"##,
387                        start_x, start_y, end_x, end_y
388                    ));
389
390                    // Event label at midpoint, offset perpendicular to line
391                    let mid_x = (start_x + end_x) / 2.0;
392                    let mid_y = (start_y + end_y) / 2.0;
393                    let perp_x = -dy / dist * 15.0;
394                    let perp_y = dx / dist * 15.0;
395
396                    svg.push_str(&format!(
397                        r##"  <text x="{:.1}" y="{:.1}" text-anchor="middle" font-family="Arial" font-size="11" fill="#0066cc">{}</text>
398"##,
399                        mid_x + perp_x, mid_y + perp_y, transition.event
400                    ));
401                }
402            }
403        }
404
405        // Draw state nodes
406        for state_name in &state_names {
407            if let Some(&(x, y)) = state_positions.get(state_name) {
408                // Circle
409                svg.push_str(&format!(
410                    r##"  <circle cx="{:.1}" cy="{:.1}" r="{}" fill="#4682b4" stroke="#2c5574" stroke-width="2"/>
411"##,
412                    x, y, node_radius
413                ));
414
415                // State name label
416                svg.push_str(&format!(
417                    "  <text x=\"{:.1}\" y=\"{:.1}\" text-anchor=\"middle\" dominant-baseline=\"middle\" font-family=\"Arial\" font-size=\"12\" fill=\"white\" font-weight=\"bold\">{}</text>\n",
418                    x, y, state_name
419                ));
420            }
421        }
422
423        // Legend
424        svg.push_str(&format!(
425            r##"  <text x="{}" y="{}" text-anchor="middle" font-family="Arial" font-size="10" fill="#666">States: circles | Transitions: arrows with event labels</text>
426"##,
427            cx, height - 15.0
428        ));
429
430        svg.push_str("</svg>\n");
431
432        // Determine output filename (convert .png to .svg if needed)
433        let output_filename = if filename.ends_with(".png") {
434            filename.replace(".png", ".svg")
435        } else {
436            filename.to_string()
437        };
438
439        // Write to file
440        if let Ok(mut file) = std::fs::File::create(&output_filename) {
441            let _ = file.write_all(svg.as_bytes());
442            println!("State machine diagram saved to {}", output_filename);
443        }
444    }
445}
446
447/// Demo robot behavior model
448pub struct RobotBehavior {
449    pub battery_level: f64,
450    pub has_obstacle: bool,
451    pub goal_reached: bool,
452    pub emergency_stop: bool,
453}
454
455impl RobotBehavior {
456    pub fn new() -> Self {
457        RobotBehavior {
458            battery_level: 100.0,
459            has_obstacle: false,
460            goal_reached: false,
461            emergency_stop: false,
462        }
463    }
464
465    pub fn can_start(&self) -> bool {
466        self.battery_level > 20.0 && !self.emergency_stop
467    }
468
469    pub fn has_battery(&self) -> bool {
470        self.battery_level > 10.0
471    }
472
473    pub fn obstacle_detected(&self) -> bool {
474        self.has_obstacle
475    }
476}
477
478impl Default for RobotBehavior {
479    fn default() -> Self {
480        Self::new()
481    }
482}
483
484/// Create a demo robot state machine
485pub fn create_robot_state_machine() -> StateMachine {
486    let mut machine = StateMachine::new("RobotController");
487
488    // Define states
489    machine.register_state(State::with_callbacks(
490        "idle",
491        Some("on_enter_idle"),
492        Some("on_exit_idle"),
493    ));
494    machine.register_state(State::with_callbacks(
495        "moving",
496        Some("on_enter_moving"),
497        Some("on_exit_moving"),
498    ));
499    machine.register_state(State::with_callbacks(
500        "avoiding",
501        Some("on_enter_avoiding"),
502        None,
503    ));
504    machine.register_state(State::with_callbacks(
505        "charging",
506        Some("on_enter_charging"),
507        None,
508    ));
509    machine.register_state(State::with_callbacks(
510        "emergency",
511        Some("on_enter_emergency"),
512        None,
513    ));
514
515    // Define transitions
516    machine.add_transition(
517        Transition::new("idle", "start", "moving")
518            .with_guard("can_start")
519            .with_action("start_motors"),
520    );
521
522    machine.add_transition(
523        Transition::new("moving", "obstacle", "avoiding").with_action("stop_motors"),
524    );
525
526    machine
527        .add_transition(Transition::new("avoiding", "clear", "moving").with_action("start_motors"));
528
529    machine.add_transition(
530        Transition::new("moving", "low_battery", "charging").with_action("save_position"),
531    );
532
533    machine
534        .add_transition(Transition::new("charging", "charged", "idle").with_action("play_sound"));
535
536    machine.add_transition(Transition::new("moving", "stop", "idle").with_action("stop_motors"));
537
538    machine.add_transition(
539        Transition::new("idle", "emergency", "emergency").with_action("send_alert"),
540    );
541
542    machine.add_transition(
543        Transition::new("moving", "emergency", "emergency").with_action("send_alert"),
544    );
545
546    machine.add_transition(Transition::new("emergency", "reset", "idle"));
547
548    machine
549}
550
551/// Run a demo of the state machine
552pub fn demo_state_machine() {
553    println!("=== Robot State Machine Demo ===\n");
554
555    // Create output directory
556    std::fs::create_dir_all("img/mission_planning").unwrap_or_default();
557
558    let mut machine = create_robot_state_machine();
559
560    // Set initial state
561    machine.set_initial_state("idle");
562
563    println!("\n{}\n", machine.generate_diagram());
564
565    // Simulate a sequence of events
566    let events = vec![
567        "start",
568        "obstacle",
569        "clear",
570        "low_battery",
571        "charged",
572        "start",
573        "emergency",
574        "reset",
575    ];
576
577    println!("=== Processing Events ===\n");
578
579    for event in events {
580        println!("Processing event: [{}]", event);
581        match machine.process(event) {
582            Ok(()) => println!("  -> Success\n"),
583            Err(e) => println!("  -> Error: {}\n", e),
584        }
585    }
586
587    // Show final state and history
588    if let Some(current) = machine.get_current_state() {
589        println!("Final state: {}", current.name);
590    }
591
592    println!("\n=== Transition History ===");
593    for (i, (from, event, to)) in machine.get_transition_history().iter().enumerate() {
594        println!("{}. {} --[{}]--> {}", i + 1, from, event, to);
595    }
596
597    // Generate visualization
598    machine.visualize("img/mission_planning/state_machine_diagram.png");
599    println!("\nState machine diagram saved to img/mission_planning/state_machine_diagram.png");
600
601    // Generate final diagram
602    println!("\n=== Final State Machine ===");
603    println!("{}", machine.generate_diagram());
604}
605
606#[cfg(test)]
607mod tests {
608    use super::*;
609
610    #[test]
611    fn test_state_creation() {
612        let state = State::new("test_state");
613        assert_eq!(state.name, "test_state");
614    }
615
616    #[test]
617    fn test_state_machine_creation() {
618        let machine = StateMachine::new("test_machine");
619        assert_eq!(machine.name, "test_machine");
620        assert!(machine.current_state.is_none());
621    }
622
623    #[test]
624    fn test_state_registration() {
625        let mut machine = StateMachine::new("test");
626        let state = State::new("idle");
627        machine.register_state(state);
628        assert!(machine.states.contains_key("idle"));
629    }
630
631    #[test]
632    fn test_transition_creation() {
633        let transition = Transition::new("idle", "start", "running")
634            .with_guard("can_start")
635            .with_action("start_motors");
636
637        assert_eq!(transition.src_state, "idle");
638        assert_eq!(transition.event, "start");
639        assert_eq!(transition.dst_state, "running");
640        assert_eq!(transition.guard, Some("can_start".to_string()));
641        assert_eq!(transition.action, Some("start_motors".to_string()));
642    }
643
644    #[test]
645    fn test_simple_transition() {
646        let mut machine = StateMachine::new("test");
647
648        machine.add_transition(Transition::new("idle", "start", "running"));
649        machine.set_initial_state("idle");
650
651        assert!(machine.process("start").is_ok());
652        assert_eq!(machine.get_current_state().unwrap().name, "running");
653    }
654
655    #[test]
656    fn test_robot_state_machine_follows_demo_transitions() {
657        let mut machine = create_robot_state_machine();
658        machine.set_initial_state("idle");
659
660        assert!(machine.process("start").is_ok());
661        assert_eq!(machine.get_current_state().unwrap().name, "moving");
662
663        assert!(machine.process("low_battery").is_ok());
664        assert_eq!(machine.get_current_state().unwrap().name, "charging");
665
666        let history = machine.get_transition_history();
667        assert_eq!(history.len(), 2);
668        assert_eq!(
669            history.last(),
670            Some(&(
671                "moving".to_string(),
672                "low_battery".to_string(),
673                "charging".to_string()
674            ))
675        );
676    }
677}
rust_robotics_control/state_machine.rs

rust_robotics_control/
state_machine.rs
