use bevy::prelude::*;
use std::collections::HashMap;
use std::hash::Hash;
#[derive(Resource, Debug)]
pub struct InputStateMachine<S, A> {
transitions: HashMap<(S, A), TransitionConfig<S>>,
global_transitions: HashMap<A, TransitionConfig<S>>,
}
impl<S: Clone + Eq + Hash, A: Clone + Eq + Hash> Default for InputStateMachine<S, A> {
fn default() -> Self {
Self::new()
}
}
impl<S: Clone + Eq + Hash, A: Clone + Eq + Hash> InputStateMachine<S, A> {
#[must_use]
pub fn new() -> Self {
Self {
transitions: HashMap::new(),
global_transitions: HashMap::new(),
}
}
pub fn add_transition(&mut self, from: S, action: A, to: S) -> &mut Self {
self.transitions.insert(
(from, action),
TransitionConfig {
target: to,
trigger: TriggerType::JustPressed,
guard: TransitionGuard::Always,
},
);
self
}
pub fn add_transition_with_trigger(
&mut self,
from: S,
action: A,
to: S,
trigger: TriggerType,
) -> &mut Self {
self.transitions.insert(
(from, action),
TransitionConfig {
target: to,
trigger,
guard: TransitionGuard::Always,
},
);
self
}
pub fn add_global_transition(&mut self, action: A, to: S) -> &mut Self {
self.global_transitions.insert(
action,
TransitionConfig {
target: to,
trigger: TriggerType::JustPressed,
guard: TransitionGuard::Always,
},
);
self
}
#[must_use]
pub fn get_transition(&self, current: &S, action: &A, trigger: TriggerType) -> Option<&S> {
if let Some(config) = self.transitions.get(&(current.clone(), action.clone()))
&& config.trigger == trigger
{
return Some(&config.target);
}
if let Some(config) = self.global_transitions.get(action)
&& config.trigger == trigger
{
return Some(&config.target);
}
None
}
}
#[derive(Debug, Clone)]
pub struct TransitionConfig<S> {
pub target: S,
pub trigger: TriggerType,
pub guard: TransitionGuard,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TriggerType {
JustPressed,
JustReleased,
Pressed,
Released,
}
#[derive(Debug, Clone)]
pub enum TransitionGuard {
Always,
Custom(String),
All(Vec<TransitionGuard>),
Any(Vec<TransitionGuard>),
}
impl TransitionGuard {
#[must_use]
pub fn always() -> Self {
Self::Always
}
#[must_use]
pub fn custom(name: impl Into<String>) -> Self {
Self::Custom(name.into())
}
}
#[derive(Component, Debug)]
pub struct InputDrivenState<S, A> {
_machine: std::marker::PhantomData<(S, A)>,
#[expect(dead_code, reason = "stored for future frame-based input processing")]
last_frame: u64,
}
impl<S, A> Default for InputDrivenState<S, A> {
fn default() -> Self {
Self {
_machine: std::marker::PhantomData,
last_frame: 0,
}
}
}
#[derive(Event, Debug, Clone)]
pub struct StateTransitionEvent<S> {
pub entity: Option<Entity>,
pub from: S,
pub to: S,
pub trigger: TriggerType,
}
#[derive(Debug)]
pub struct StateMachineBuilder<S, A> {
machine: InputStateMachine<S, A>,
}
impl<S: Clone + Eq + Hash, A: Clone + Eq + Hash> StateMachineBuilder<S, A> {
#[must_use]
pub fn new() -> Self {
Self {
machine: InputStateMachine::new(),
}
}
#[must_use]
pub fn on(mut self, from: S, action: A, to: S) -> Self {
self.machine.add_transition(from, action, to);
self
}
#[must_use]
pub fn on_trigger(mut self, from: S, action: A, to: S, trigger: TriggerType) -> Self {
self.machine
.add_transition_with_trigger(from, action, to, trigger);
self
}
#[must_use]
pub fn on_any(mut self, action: A, to: S) -> Self {
self.machine.add_global_transition(action, to);
self
}
#[must_use]
pub fn build(self) -> InputStateMachine<S, A> {
self.machine
}
}
impl<S: Clone + Eq + Hash, A: Clone + Eq + Hash> Default for StateMachineBuilder<S, A> {
fn default() -> Self {
Self::new()
}
}
#[derive(Debug, Default)]
pub struct StateGraph<S> {
pub states: Vec<S>,
pub edges: Vec<(usize, usize, String)>,
}
impl<S: Clone + Eq + Hash> StateGraph<S> {
#[must_use]
pub fn from_machine<A: Clone + Eq + Hash + std::fmt::Debug>(
machine: &InputStateMachine<S, A>,
) -> Self {
let mut graph = Self {
states: Vec::new(),
edges: Vec::new(),
};
let mut state_indices: HashMap<S, usize> = HashMap::new();
for ((from, _action), config) in &machine.transitions {
if !state_indices.contains_key(from) {
state_indices.insert(from.clone(), graph.states.len());
graph.states.push(from.clone());
}
if !state_indices.contains_key(&config.target) {
state_indices.insert(config.target.clone(), graph.states.len());
graph.states.push(config.target.clone());
}
}
for ((from, action), config) in &machine.transitions {
let from_idx = state_indices[from];
let to_idx = state_indices[&config.target];
graph.edges.push((from_idx, to_idx, format!("{action:?}")));
}
graph
}
}
#[derive(Component, Debug)]
pub struct TimedState<S> {
pub current: S,
pub timer: f32,
pub next: Option<S>,
}
impl<S: Clone> TimedState<S> {
#[must_use]
pub fn new(initial: S) -> Self {
Self {
current: initial,
timer: 0.0,
next: None,
}
}
pub fn transition_after(&mut self, duration: f32, to: S) {
self.timer = duration;
self.next = Some(to);
}
pub fn update(&mut self, delta: f32) -> Option<S> {
if self.next.is_some() {
self.timer -= delta;
if self.timer <= 0.0
&& let Some(next) = self.next.take()
{
self.current = next.clone();
return Some(next);
}
}
None
}
}
#[derive(SystemSet, Debug, Clone, PartialEq, Eq, Hash)]
pub enum StateMachineSet {
ProcessInput,
Transition,
PostTransition,
}
pub struct StateMachinePlugin;
impl Plugin for StateMachinePlugin {
fn build(&self, app: &mut App) {
app.configure_sets(
Update,
(
StateMachineSet::ProcessInput,
StateMachineSet::Transition,
StateMachineSet::PostTransition,
)
.chain(),
);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[derive(States, Clone, Eq, PartialEq, Debug, Hash, Default)]
enum TestState {
#[default]
Idle,
Running,
Jumping,
}
#[derive(Clone, Eq, PartialEq, Debug, Hash)]
enum TestAction {
Move,
Jump,
}
#[test]
fn test_state_machine_basic() {
let mut machine = InputStateMachine::<TestState, TestAction>::new();
machine.add_transition(TestState::Idle, TestAction::Jump, TestState::Jumping);
machine.add_transition(TestState::Idle, TestAction::Move, TestState::Running);
let target = machine.get_transition(
&TestState::Idle,
&TestAction::Jump,
TriggerType::JustPressed,
);
assert_eq!(target, Some(&TestState::Jumping));
let target = machine.get_transition(
&TestState::Running,
&TestAction::Jump,
TriggerType::JustPressed,
);
assert_eq!(target, None); }
#[test]
fn test_global_transition() {
let mut machine = InputStateMachine::<TestState, TestAction>::new();
machine.add_global_transition(TestAction::Jump, TestState::Jumping);
let target = machine.get_transition(
&TestState::Running, &TestAction::Jump,
TriggerType::JustPressed,
);
assert_eq!(target, Some(&TestState::Jumping));
}
#[test]
fn test_builder() {
let machine = StateMachineBuilder::<TestState, TestAction>::new()
.on(TestState::Idle, TestAction::Jump, TestState::Jumping)
.on(TestState::Idle, TestAction::Move, TestState::Running)
.on_any(TestAction::Jump, TestState::Jumping)
.build();
assert!(machine.transitions.len() >= 2);
}
#[test]
fn test_timed_state() {
let mut timed = TimedState::new(TestState::Jumping);
timed.transition_after(0.5, TestState::Idle);
assert!(timed.update(0.3).is_none());
let result = timed.update(0.3);
assert_eq!(result, Some(TestState::Idle));
assert_eq!(timed.current, TestState::Idle);
}
#[test]
fn test_state_graph() {
let machine = StateMachineBuilder::new()
.on(TestState::Idle, TestAction::Jump, TestState::Jumping)
.on(TestState::Jumping, TestAction::Move, TestState::Running)
.build();
let graph = StateGraph::from_machine(&machine);
assert_eq!(graph.states.len(), 3);
assert_eq!(graph.edges.len(), 2);
}
}