use bevy_app::{App, MainScheduleOrder, Plugin, PreUpdate, Startup, SubApp};
use bevy_ecs::{
event::Events,
schedule::{IntoSystemConfigs, ScheduleLabel},
world::FromWorld,
};
use bevy_utils::{tracing::warn, warn_once};
use crate::state::{
setup_state_transitions_in_world, ComputedStates, FreelyMutableState, NextState, State,
StateTransition, StateTransitionEvent, StateTransitionSteps, States, SubStates,
};
use crate::state_scoped::clear_state_scoped_entities;
pub trait AppExtStates {
fn init_state<S: FreelyMutableState + FromWorld>(&mut self) -> &mut Self;
fn insert_state<S: FreelyMutableState>(&mut self, state: S) -> &mut Self;
fn add_computed_state<S: ComputedStates>(&mut self) -> &mut Self;
fn add_sub_state<S: SubStates>(&mut self) -> &mut Self;
fn enable_state_scoped_entities<S: States>(&mut self) -> &mut Self;
}
fn warn_if_no_states_plugin_installed(app: &SubApp) {
if !app.is_plugin_added::<StatesPlugin>() {
warn_once!("States were added to the app, but `StatesPlugin` is not installed.");
}
}
impl AppExtStates for SubApp {
fn init_state<S: FreelyMutableState + FromWorld>(&mut self) -> &mut Self {
warn_if_no_states_plugin_installed(self);
if !self.world().contains_resource::<State<S>>() {
self.init_resource::<State<S>>()
.init_resource::<NextState<S>>()
.add_event::<StateTransitionEvent<S>>();
let schedule = self.get_schedule_mut(StateTransition).unwrap();
S::register_state(schedule);
let state = self.world().resource::<State<S>>().get().clone();
self.world_mut().send_event(StateTransitionEvent {
exited: None,
entered: Some(state),
});
} else {
let name = std::any::type_name::<S>();
warn!("State {} is already initialized.", name);
}
self
}
fn insert_state<S: FreelyMutableState>(&mut self, state: S) -> &mut Self {
warn_if_no_states_plugin_installed(self);
if !self.world().contains_resource::<State<S>>() {
self.insert_resource::<State<S>>(State::new(state.clone()))
.init_resource::<NextState<S>>()
.add_event::<StateTransitionEvent<S>>();
let schedule = self.get_schedule_mut(StateTransition).unwrap();
S::register_state(schedule);
self.world_mut().send_event(StateTransitionEvent {
exited: None,
entered: Some(state),
});
} else {
self.insert_resource::<State<S>>(State::new(state.clone()));
self.world_mut()
.resource_mut::<Events<StateTransitionEvent<S>>>()
.clear();
self.world_mut().send_event(StateTransitionEvent {
exited: None,
entered: Some(state),
});
}
self
}
fn add_computed_state<S: ComputedStates>(&mut self) -> &mut Self {
warn_if_no_states_plugin_installed(self);
if !self
.world()
.contains_resource::<Events<StateTransitionEvent<S>>>()
{
self.add_event::<StateTransitionEvent<S>>();
let schedule = self.get_schedule_mut(StateTransition).unwrap();
S::register_computed_state_systems(schedule);
let state = self
.world()
.get_resource::<State<S>>()
.map(|s| s.get().clone());
self.world_mut().send_event(StateTransitionEvent {
exited: None,
entered: state,
});
} else {
let name = std::any::type_name::<S>();
warn!("Computed state {} is already initialized.", name);
}
self
}
fn add_sub_state<S: SubStates>(&mut self) -> &mut Self {
warn_if_no_states_plugin_installed(self);
if !self
.world()
.contains_resource::<Events<StateTransitionEvent<S>>>()
{
self.init_resource::<NextState<S>>();
self.add_event::<StateTransitionEvent<S>>();
let schedule = self.get_schedule_mut(StateTransition).unwrap();
S::register_sub_state_systems(schedule);
let state = self
.world()
.get_resource::<State<S>>()
.map(|s| s.get().clone());
self.world_mut().send_event(StateTransitionEvent {
exited: None,
entered: state,
});
} else {
let name = std::any::type_name::<S>();
warn!("Sub state {} is already initialized.", name);
}
self
}
fn enable_state_scoped_entities<S: States>(&mut self) -> &mut Self {
if !self
.world()
.contains_resource::<Events<StateTransitionEvent<S>>>()
{
let name = std::any::type_name::<S>();
warn!("State scoped entities are enabled for state `{}`, but the state isn't installed in the app!", name);
}
self.add_systems(
StateTransition,
clear_state_scoped_entities::<S>.in_set(StateTransitionSteps::ExitSchedules),
)
}
}
impl AppExtStates for App {
fn init_state<S: FreelyMutableState + FromWorld>(&mut self) -> &mut Self {
self.main_mut().init_state::<S>();
self
}
fn insert_state<S: FreelyMutableState>(&mut self, state: S) -> &mut Self {
self.main_mut().insert_state::<S>(state);
self
}
fn add_computed_state<S: ComputedStates>(&mut self) -> &mut Self {
self.main_mut().add_computed_state::<S>();
self
}
fn add_sub_state<S: SubStates>(&mut self) -> &mut Self {
self.main_mut().add_sub_state::<S>();
self
}
fn enable_state_scoped_entities<S: States>(&mut self) -> &mut Self {
self.main_mut().enable_state_scoped_entities::<S>();
self
}
}
pub struct StatesPlugin;
impl Plugin for StatesPlugin {
fn build(&self, app: &mut App) {
let mut schedule = app.world_mut().resource_mut::<MainScheduleOrder>();
schedule.insert_after(PreUpdate, StateTransition);
setup_state_transitions_in_world(app.world_mut(), Some(Startup.intern()));
}
}
#[cfg(test)]
mod tests {
use crate::{
self as bevy_state,
app::StatesPlugin,
state::{State, StateTransition, StateTransitionEvent},
};
use bevy_app::App;
use bevy_ecs::event::Events;
use bevy_state_macros::States;
use super::AppExtStates;
#[derive(States, Default, PartialEq, Eq, Hash, Debug, Clone)]
enum TestState {
#[default]
A,
B,
C,
}
#[test]
fn insert_state_can_overwrite_init_state() {
let mut app = App::new();
app.add_plugins(StatesPlugin);
app.init_state::<TestState>();
app.insert_state(TestState::B);
let world = app.world_mut();
world.run_schedule(StateTransition);
assert_eq!(world.resource::<State<TestState>>().0, TestState::B);
let events = world.resource::<Events<StateTransitionEvent<TestState>>>();
assert_eq!(events.len(), 1);
let mut reader = events.get_reader();
let last = reader.read(events).last().unwrap();
assert_eq!(last.exited, None);
assert_eq!(last.entered, Some(TestState::B));
}
#[test]
fn insert_state_can_overwrite_insert_state() {
let mut app = App::new();
app.add_plugins(StatesPlugin);
app.insert_state(TestState::B);
app.insert_state(TestState::C);
let world = app.world_mut();
world.run_schedule(StateTransition);
assert_eq!(world.resource::<State<TestState>>().0, TestState::C);
let events = world.resource::<Events<StateTransitionEvent<TestState>>>();
assert_eq!(events.len(), 1);
let mut reader = events.get_reader();
let last = reader.read(events).last().unwrap();
assert_eq!(last.exited, None);
assert_eq!(last.entered, Some(TestState::C));
}
}