Skip to main content

qubit_state_machine/
state_machine.rs

1/*******************************************************************************
2 *
3 *    Copyright (c) 2026 Haixing Hu.
4 *
5 *    SPDX-License-Identifier: Apache-2.0
6 *
7 *    Licensed under the Apache License, Version 2.0.
8 *
9 ******************************************************************************/
10//! Immutable finite state machine rules and CAS-backed event triggering.
11
12use std::collections::{
13    HashMap,
14    HashSet,
15};
16use std::fmt::Debug;
17use std::hash::Hash;
18
19use qubit_atomic::AtomicRef;
20use qubit_cas::{
21    CasDecision,
22    CasError,
23    CasExecutor,
24    CasSuccess,
25};
26
27use crate::{
28    StateMachineBuilder,
29    StateMachineError,
30    StateMachineResult,
31    Transition,
32};
33
34/// Immutable finite state machine rules.
35///
36/// `S` is the state type and `E` is the event type. Both should usually be
37/// small enum-like types. The state machine itself is immutable and can be
38/// shared across threads; mutable current state is kept in [`AtomicRef`] and
39/// updated through [`qubit_cas::CasExecutor`].
40///
41/// # Common usage
42///
43/// Define the valid states and events, build an immutable transition table, and
44/// keep each object's current state in an [`AtomicRef`].
45///
46/// ```
47/// use qubit_state_machine::{AtomicRef, StateMachine};
48///
49/// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
50/// enum JobState {
51///     Queued,
52///     Running,
53///     Succeeded,
54///     Failed,
55/// }
56///
57/// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
58/// enum JobEvent {
59///     Start,
60///     Complete,
61///     Fail,
62/// }
63///
64/// fn create_job_machine() -> StateMachine<JobState, JobEvent> {
65///     StateMachine::builder()
66///         .add_states(&[
67///             JobState::Queued,
68///             JobState::Running,
69///             JobState::Succeeded,
70///             JobState::Failed,
71///         ])
72///         .set_initial_state(JobState::Queued)
73///         .set_final_states(&[JobState::Succeeded, JobState::Failed])
74///         .add_transition(JobState::Queued, JobEvent::Start, JobState::Running)
75///         .add_transition(JobState::Running, JobEvent::Complete, JobState::Succeeded)
76///         .add_transition(JobState::Running, JobEvent::Fail, JobState::Failed)
77///         .build()
78///         .expect("job state machine should be valid")
79/// }
80///
81/// let machine = create_job_machine();
82/// let state = AtomicRef::from_value(JobState::Queued);
83///
84/// assert_eq!(machine.trigger(&state, JobEvent::Start).unwrap(), JobState::Running);
85/// assert_eq!(*state.load(), JobState::Running);
86///
87/// assert!(machine.try_trigger(&state, JobEvent::Complete));
88/// assert_eq!(*state.load(), JobState::Succeeded);
89/// ```
90#[derive(Debug, Clone)]
91pub struct StateMachine<S, E>
92where
93    S: Copy + Eq + Hash + Debug + 'static,
94    E: Copy + Eq + Hash + Debug + 'static,
95{
96    states: HashSet<S>,
97    initial_states: HashSet<S>,
98    final_states: HashSet<S>,
99    transitions: HashSet<Transition<S, E>>,
100    transition_map: HashMap<(S, E), S>,
101    cas_executor: CasExecutor<S, StateMachineError<S, E>>,
102}
103
104impl<S, E> StateMachine<S, E>
105where
106    S: Copy + Eq + Hash + Debug + 'static,
107    E: Copy + Eq + Hash + Debug + 'static,
108{
109    /// Creates a builder for immutable state machine rules.
110    ///
111    /// # Returns
112    /// A new empty [`StateMachineBuilder`].
113    ///
114    /// # Examples
115    ///
116    /// ```
117    /// use qubit_state_machine::StateMachine;
118    ///
119    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
120    /// enum State {
121    ///     New,
122    /// }
123    ///
124    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
125    /// enum Event {
126    ///     Start,
127    /// }
128    ///
129    /// let machine = StateMachine::<State, Event>::builder()
130    ///     .add_state(State::New)
131    ///     .set_initial_state(State::New)
132    ///     .build()
133    ///     .expect("single-state machine should build");
134    /// assert!(machine.contains_state(State::New));
135    /// ```
136    pub fn builder() -> StateMachineBuilder<S, E> {
137        StateMachineBuilder::new()
138    }
139
140    /// Creates a state machine from validated builder parts.
141    ///
142    /// # Parameters
143    /// - `builder`: Builder containing validated states and terminal markers.
144    /// - `transitions`: Validated transition set.
145    /// - `transition_map`: Lookup table keyed by `(source, event)`.
146    ///
147    /// # Returns
148    /// An immutable state machine.
149    ///
150    /// This constructor does not validate input. Rule validation belongs to
151    /// [`StateMachineBuilder::build`].
152    pub(crate) fn new(
153        builder: StateMachineBuilder<S, E>,
154        transitions: HashSet<Transition<S, E>>,
155        transition_map: HashMap<(S, E), S>,
156    ) -> Self {
157        Self {
158            states: builder.states,
159            initial_states: builder.initial_states,
160            final_states: builder.final_states,
161            transitions,
162            transition_map,
163            cas_executor: CasExecutor::latency_first(),
164        }
165    }
166
167    /// Returns all registered states.
168    ///
169    /// # Returns
170    /// An immutable view of the registered state set.
171    ///
172    /// # Examples
173    ///
174    /// ```
175    /// # use qubit_state_machine::StateMachine;
176    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
177    /// # enum State { New, Running }
178    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
179    /// # enum Event { Start }
180    /// # let machine = StateMachine::builder()
181    /// #     .add_states(&[State::New, State::Running])
182    /// #     .add_transition(State::New, Event::Start, State::Running)
183    /// #     .build()
184    /// #     .expect("rules should build");
185    /// assert!(machine.states().contains(&State::New));
186    /// assert_eq!(machine.states().len(), 2);
187    /// ```
188    pub const fn states(&self) -> &HashSet<S> {
189        &self.states
190    }
191
192    /// Returns all configured initial states.
193    ///
194    /// # Returns
195    /// An immutable view of the initial state set.
196    ///
197    /// # Examples
198    ///
199    /// ```
200    /// # use qubit_state_machine::StateMachine;
201    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
202    /// # enum State { New, Running }
203    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
204    /// # enum Event { Start }
205    /// # let machine = StateMachine::builder()
206    /// #     .add_states(&[State::New, State::Running])
207    /// #     .set_initial_state(State::New)
208    /// #     .add_transition(State::New, Event::Start, State::Running)
209    /// #     .build()
210    /// #     .expect("rules should build");
211    /// assert!(machine.initial_states().contains(&State::New));
212    /// ```
213    pub const fn initial_states(&self) -> &HashSet<S> {
214        &self.initial_states
215    }
216
217    /// Returns all configured final states.
218    ///
219    /// # Returns
220    /// An immutable view of the final state set.
221    ///
222    /// # Examples
223    ///
224    /// ```
225    /// # use qubit_state_machine::StateMachine;
226    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
227    /// # enum State { New, Done }
228    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
229    /// # enum Event { Finish }
230    /// # let machine = StateMachine::builder()
231    /// #     .add_states(&[State::New, State::Done])
232    /// #     .set_final_state(State::Done)
233    /// #     .add_transition(State::New, Event::Finish, State::Done)
234    /// #     .build()
235    /// #     .expect("rules should build");
236    /// assert!(machine.final_states().contains(&State::Done));
237    /// ```
238    pub const fn final_states(&self) -> &HashSet<S> {
239        &self.final_states
240    }
241
242    /// Returns all registered transitions.
243    ///
244    /// # Returns
245    /// An immutable view of the transition set.
246    ///
247    /// # Examples
248    ///
249    /// ```
250    /// use qubit_state_machine::{StateMachine, Transition};
251    ///
252    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
253    /// enum State {
254    ///     New,
255    ///     Running,
256    /// }
257    ///
258    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
259    /// enum Event {
260    ///     Start,
261    /// }
262    ///
263    /// let machine = StateMachine::builder()
264    ///     .add_states(&[State::New, State::Running])
265    ///     .add_transition(State::New, Event::Start, State::Running)
266    ///     .build()
267    ///     .expect("rules should build");
268    ///
269    /// assert!(machine
270    ///     .transitions()
271    ///     .contains(&Transition::new(State::New, Event::Start, State::Running)));
272    /// ```
273    pub const fn transitions(&self) -> &HashSet<Transition<S, E>> {
274        &self.transitions
275    }
276
277    /// Tests whether a state is registered in this state machine.
278    ///
279    /// # Parameters
280    /// - `state`: State to test.
281    ///
282    /// # Returns
283    /// `true` if the state is registered.
284    ///
285    /// # Examples
286    ///
287    /// ```
288    /// # use qubit_state_machine::StateMachine;
289    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
290    /// # enum State { New, Running, Detached }
291    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
292    /// # enum Event { Start }
293    /// # let machine = StateMachine::builder()
294    /// #     .add_states(&[State::New, State::Running])
295    /// #     .add_transition(State::New, Event::Start, State::Running)
296    /// #     .build()
297    /// #     .expect("rules should build");
298    /// assert!(machine.contains_state(State::Running));
299    /// assert!(!machine.contains_state(State::Detached));
300    /// ```
301    pub fn contains_state(&self, state: S) -> bool {
302        self.states.contains(&state)
303    }
304
305    /// Tests whether a state is configured as an initial state.
306    ///
307    /// # Parameters
308    /// - `state`: State to test.
309    ///
310    /// # Returns
311    /// `true` if the state is an initial state.
312    ///
313    /// # Examples
314    ///
315    /// ```
316    /// # use qubit_state_machine::StateMachine;
317    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
318    /// # enum State { New, Running }
319    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
320    /// # enum Event { Start }
321    /// # let machine = StateMachine::builder()
322    /// #     .add_states(&[State::New, State::Running])
323    /// #     .set_initial_state(State::New)
324    /// #     .add_transition(State::New, Event::Start, State::Running)
325    /// #     .build()
326    /// #     .expect("rules should build");
327    /// assert!(machine.is_initial_state(State::New));
328    /// assert!(!machine.is_initial_state(State::Running));
329    /// ```
330    pub fn is_initial_state(&self, state: S) -> bool {
331        self.initial_states.contains(&state)
332    }
333
334    /// Tests whether a state is configured as a final state.
335    ///
336    /// # Parameters
337    /// - `state`: State to test.
338    ///
339    /// # Returns
340    /// `true` if the state is a final state.
341    ///
342    /// # Examples
343    ///
344    /// ```
345    /// # use qubit_state_machine::StateMachine;
346    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
347    /// # enum State { Running, Done }
348    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
349    /// # enum Event { Finish }
350    /// # let machine = StateMachine::builder()
351    /// #     .add_states(&[State::Running, State::Done])
352    /// #     .set_final_state(State::Done)
353    /// #     .add_transition(State::Running, Event::Finish, State::Done)
354    /// #     .build()
355    /// #     .expect("rules should build");
356    /// assert!(machine.is_final_state(State::Done));
357    /// assert!(!machine.is_final_state(State::Running));
358    /// ```
359    pub fn is_final_state(&self, state: S) -> bool {
360        self.final_states.contains(&state)
361    }
362
363    /// Looks up the target state for a source state and event.
364    ///
365    /// This method only queries rules; it does not modify any current-state
366    /// storage.
367    ///
368    /// # Parameters
369    /// - `source`: Source state.
370    /// - `event`: Event to apply.
371    ///
372    /// # Returns
373    /// `Some(target)` if a transition exists, or `None` otherwise.
374    ///
375    /// # Examples
376    ///
377    /// ```
378    /// # use qubit_state_machine::StateMachine;
379    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
380    /// # enum State { New, Running }
381    /// # #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
382    /// # enum Event { Start, Finish }
383    /// # let machine = StateMachine::builder()
384    /// #     .add_states(&[State::New, State::Running])
385    /// #     .add_transition(State::New, Event::Start, State::Running)
386    /// #     .build()
387    /// #     .expect("rules should build");
388    /// assert_eq!(
389    ///     machine.transition_target(State::New, Event::Start),
390    ///     Some(State::Running),
391    /// );
392    /// assert_eq!(machine.transition_target(State::New, Event::Finish), None);
393    /// ```
394    pub fn transition_target(&self, source: S, event: E) -> Option<S> {
395        self.transition_map.get(&(source, event)).copied()
396    }
397
398    /// Triggers an event and updates the provided atomic state reference.
399    ///
400    /// # Parameters
401    /// - `state`: Current state atomic reference.
402    /// - `event`: Event to apply.
403    ///
404    /// # Returns
405    /// The new state after a successful transition.
406    ///
407    /// # Errors
408    /// Returns [`StateMachineError::UnknownState`] when the current state is not
409    /// registered. Returns [`StateMachineError::UnknownTransition`] when the
410    /// current state is registered but has no transition for `event`.
411    ///
412    /// # Examples
413    ///
414    /// ```
415    /// use qubit_state_machine::{AtomicRef, StateMachine};
416    ///
417    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
418    /// enum State {
419    ///     New,
420    ///     Running,
421    /// }
422    ///
423    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
424    /// enum Event {
425    ///     Start,
426    /// }
427    ///
428    /// let machine = StateMachine::builder()
429    ///     .add_states(&[State::New, State::Running])
430    ///     .add_transition(State::New, Event::Start, State::Running)
431    ///     .build()
432    ///     .expect("rules should build");
433    /// let state = AtomicRef::from_value(State::New);
434    ///
435    /// assert_eq!(machine.trigger(&state, Event::Start).unwrap(), State::Running);
436    /// assert_eq!(*state.load(), State::Running);
437    /// ```
438    pub fn trigger(&self, state: &AtomicRef<S>, event: E) -> StateMachineResult<S, E> {
439        let (_, new_state) = self.change_state(state, event)?;
440        Ok(new_state)
441    }
442
443    /// Triggers an event, updates the atomic state, and invokes a success callback.
444    ///
445    /// The callback runs after the CAS update has succeeded.
446    ///
447    /// # Parameters
448    /// - `state`: Current state atomic reference.
449    /// - `event`: Event to apply.
450    /// - `on_success`: Callback receiving `(old_state, new_state)`.
451    ///
452    /// # Returns
453    /// The new state after a successful transition.
454    ///
455    /// # Errors
456    /// Returns the same errors as [`StateMachine::trigger`]. The callback is not
457    /// invoked when the transition fails.
458    ///
459    /// # Examples
460    ///
461    /// ```
462    /// use qubit_state_machine::{AtomicRef, StateMachine};
463    ///
464    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
465    /// enum State {
466    ///     New,
467    ///     Running,
468    /// }
469    ///
470    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
471    /// enum Event {
472    ///     Start,
473    /// }
474    ///
475    /// let machine = StateMachine::builder()
476    ///     .add_states(&[State::New, State::Running])
477    ///     .add_transition(State::New, Event::Start, State::Running)
478    ///     .build()
479    ///     .expect("rules should build");
480    /// let state = AtomicRef::from_value(State::New);
481    /// let mut observed = None;
482    ///
483    /// let next = machine
484    ///     .trigger_with(&state, Event::Start, |old_state, new_state| {
485    ///         observed = Some((old_state, new_state));
486    ///     })
487    ///     .expect("start should be valid");
488    ///
489    /// assert_eq!(next, State::Running);
490    /// assert_eq!(observed, Some((State::New, State::Running)));
491    /// ```
492    pub fn trigger_with<F>(
493        &self,
494        state: &AtomicRef<S>,
495        event: E,
496        on_success: F,
497    ) -> StateMachineResult<S, E>
498    where
499        F: FnOnce(S, S),
500    {
501        let (old_state, new_state) = self.change_state(state, event)?;
502        on_success(old_state, new_state);
503        Ok(new_state)
504    }
505
506    /// Attempts to trigger an event without returning error details.
507    ///
508    /// # Parameters
509    /// - `state`: Current state atomic reference.
510    /// - `event`: Event to apply.
511    ///
512    /// # Returns
513    /// `true` if the state changed successfully; `false` if the transition was
514    /// invalid.
515    ///
516    /// # Examples
517    ///
518    /// ```
519    /// use qubit_state_machine::{AtomicRef, StateMachine};
520    ///
521    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
522    /// enum State {
523    ///     New,
524    ///     Running,
525    /// }
526    ///
527    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
528    /// enum Event {
529    ///     Start,
530    ///     Finish,
531    /// }
532    ///
533    /// let machine = StateMachine::builder()
534    ///     .add_states(&[State::New, State::Running])
535    ///     .add_transition(State::New, Event::Start, State::Running)
536    ///     .build()
537    ///     .expect("rules should build");
538    /// let state = AtomicRef::from_value(State::New);
539    ///
540    /// assert!(!machine.try_trigger(&state, Event::Finish));
541    /// assert_eq!(*state.load(), State::New);
542    /// assert!(machine.try_trigger(&state, Event::Start));
543    /// ```
544    pub fn try_trigger(&self, state: &AtomicRef<S>, event: E) -> bool {
545        self.trigger(state, event).is_ok()
546    }
547
548    /// Attempts to trigger an event and invokes a callback only on success.
549    ///
550    /// # Parameters
551    /// - `state`: Current state atomic reference.
552    /// - `event`: Event to apply.
553    /// - `on_success`: Callback receiving `(old_state, new_state)`.
554    ///
555    /// # Returns
556    /// `true` if the state changed successfully; `false` if the transition was
557    /// invalid. The callback is skipped when this method returns `false`.
558    ///
559    /// # Examples
560    ///
561    /// ```
562    /// use qubit_state_machine::{AtomicRef, StateMachine};
563    ///
564    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
565    /// enum State {
566    ///     New,
567    ///     Running,
568    /// }
569    ///
570    /// #[derive(Debug, Clone, Copy, Eq, PartialEq, Hash)]
571    /// enum Event {
572    ///     Start,
573    ///     Finish,
574    /// }
575    ///
576    /// let machine = StateMachine::builder()
577    ///     .add_states(&[State::New, State::Running])
578    ///     .add_transition(State::New, Event::Start, State::Running)
579    ///     .build()
580    ///     .expect("rules should build");
581    /// let state = AtomicRef::from_value(State::New);
582    /// let mut callback_count = 0;
583    ///
584    /// assert!(!machine.try_trigger_with(&state, Event::Finish, |_, _| {
585    ///     callback_count += 1;
586    /// }));
587    /// assert_eq!(callback_count, 0);
588    ///
589    /// assert!(machine.try_trigger_with(&state, Event::Start, |_, _| {
590    ///     callback_count += 1;
591    /// }));
592    /// assert_eq!(callback_count, 1);
593    /// ```
594    pub fn try_trigger_with<F>(&self, state: &AtomicRef<S>, event: E, on_success: F) -> bool
595    where
596        F: FnOnce(S, S),
597    {
598        self.trigger_with(state, event, on_success).is_ok()
599    }
600
601    /// Applies a transition through the CAS executor.
602    ///
603    /// # Parameters
604    /// - `state`: Current state atomic reference.
605    /// - `event`: Event to apply.
606    ///
607    /// # Returns
608    /// The old and new state when the transition succeeds.
609    ///
610    /// # Errors
611    /// Returns a runtime state machine error when no valid next state exists.
612    fn change_state(
613        &self,
614        state: &AtomicRef<S>,
615        event: E,
616    ) -> Result<(S, S), StateMachineError<S, E>> {
617        let outcome = self.cas_executor.execute(state, |current_state: &S| {
618            match self.next_state(*current_state, event) {
619                Ok(new_state) => CasDecision::update(new_state, new_state),
620                Err(error) => CasDecision::abort(error),
621            }
622        });
623        match outcome.into_result() {
624            Ok(success) => Ok(Self::state_change_from_success(success)),
625            Err(error) => Err(Self::state_error_from_cas_error(error)),
626        }
627    }
628
629    /// Resolves the next state for the current state and event.
630    ///
631    /// # Parameters
632    /// - `current_state`: State currently stored by the atomic reference.
633    /// - `event`: Event to apply.
634    ///
635    /// # Returns
636    /// The target state when a transition exists.
637    ///
638    /// # Errors
639    /// Returns an unknown-state error before checking transitions if the current
640    /// state is not registered. Returns an unknown-transition error if no rule
641    /// exists for the `(current_state, event)` pair.
642    fn next_state(&self, current_state: S, event: E) -> Result<S, StateMachineError<S, E>> {
643        if !self.contains_state(current_state) {
644            return Err(StateMachineError::UnknownState {
645                state: current_state,
646            });
647        }
648        self.transition_target(current_state, event)
649            .ok_or(StateMachineError::UnknownTransition {
650                source: current_state,
651                event,
652            })
653    }
654
655    /// Extracts old and new states from a successful CAS transition.
656    ///
657    /// # Parameters
658    /// - `success`: Successful CAS result returned by the executor.
659    ///
660    /// # Returns
661    /// The old state and current state after CAS completion.
662    fn state_change_from_success(success: CasSuccess<S, S>) -> (S, S) {
663        match success {
664            CasSuccess::Updated {
665                previous, current, ..
666            } => (*previous, *current),
667            CasSuccess::Finished { current, .. } => (*current, *current),
668        }
669    }
670
671    /// Maps terminal CAS failures into state machine errors.
672    ///
673    /// # Parameters
674    /// - `error`: Terminal CAS error returned by the executor.
675    ///
676    /// # Returns
677    /// The business state machine error when the operation aborted, or a CAS
678    /// conflict error when retry limits were exhausted by compare-and-swap
679    /// conflicts.
680    fn state_error_from_cas_error(
681        error: CasError<S, StateMachineError<S, E>>,
682    ) -> StateMachineError<S, E> {
683        match error.error() {
684            Some(error) => *error,
685            None => StateMachineError::CasConflict {
686                attempts: error.attempts(),
687            },
688        }
689    }
690}