use bevy_app::{App, CoreStage};
use bevy_ecs::schedule::{IntoSystemDescriptor, StageLabel, StateData, SystemSet};
use crate::{Hook, StackStage};
pub trait AppStateStackExt {
fn add_state_stack<T: StateData + StageLabel>(&mut self, stack: T) -> &mut App;
fn add_system_on_enter<T, Params>(
&mut self,
state: T,
system: impl IntoSystemDescriptor<Params>,
) -> &mut Self
where
T: StateData + StageLabel;
fn add_system_on_exit<T, Params>(
&mut self,
state: T,
system: impl IntoSystemDescriptor<Params>,
) -> &mut Self
where
T: StateData + StageLabel;
fn add_system_on_pause<T, Params>(
&mut self,
state: T,
system: impl IntoSystemDescriptor<Params>,
) -> &mut Self
where
T: StateData + StageLabel;
fn add_system_on_resume<T, Params>(
&mut self,
state: T,
system: impl IntoSystemDescriptor<Params>,
) -> &mut Self
where
T: StateData + StageLabel;
fn add_system_on_update<T, Params>(
&mut self,
state: T,
system: impl IntoSystemDescriptor<Params>,
) -> &mut Self
where
T: StateData + StageLabel;
fn add_system_set_on_update<T>(&mut self, state: T, system: SystemSet) -> &mut Self
where
T: StateData + StageLabel;
}
fn into_label<T: StageLabel>(s: &T) -> &'static str {
s.as_str().split("::").next().unwrap()
}
impl AppStateStackExt for App {
fn add_state_stack<T: StateData + StageLabel>(&mut self, stack: T) -> &mut App {
self.add_stage_before(
CoreStage::Update,
into_label(&stack),
StackStage::new(stack),
)
}
fn add_system_on_exit<T, Params>(
&mut self,
state: T,
system: impl IntoSystemDescriptor<Params>,
) -> &mut Self
where
T: StateData + StageLabel,
{
let stage = self
.schedule
.get_stage_mut::<StackStage<T>>(&into_label(&state))
.expect("No state stack with that type! Use add_state_stack to add it!");
stage.add_system(Hook::Exit, state, system);
self
}
fn add_system_on_enter<T, Params>(
&mut self,
state: T,
system: impl IntoSystemDescriptor<Params>,
) -> &mut Self
where
T: StateData + StageLabel,
{
let stage = self
.schedule
.get_stage_mut::<StackStage<T>>(&into_label(&state))
.expect("No state stack with that type! Use add_state_stack to add it!");
stage.add_system(Hook::Enter, state, system);
self
}
fn add_system_on_pause<T, Params>(
&mut self,
state: T,
system: impl IntoSystemDescriptor<Params>,
) -> &mut Self
where
T: StateData + StageLabel,
{
let stage = self
.schedule
.get_stage_mut::<StackStage<T>>(&into_label(&state))
.expect("No state stack with that type! Use add_state_stack to add it!");
stage.add_system(Hook::Pause, state, system);
self
}
fn add_system_on_resume<T, Params>(
&mut self,
state: T,
system: impl IntoSystemDescriptor<Params>,
) -> &mut Self
where
T: StateData + StageLabel,
{
let stage = self
.schedule
.get_stage_mut::<StackStage<T>>(&into_label(&state))
.expect("No state stack with that type! Use add_state_stack to add it!");
stage.add_system(Hook::Resume, state, system);
self
}
fn add_system_on_update<T, Params>(
&mut self,
state: T,
system: impl IntoSystemDescriptor<Params>,
) -> &mut Self
where
T: StateData + StageLabel,
{
self.add_system_set(
SystemSet::new()
.with_run_criteria(crate::in_state::<T>(state))
.with_system(system),
);
self
}
fn add_system_set_on_update<T>(&mut self, state: T, system_set: SystemSet) -> &mut Self
where
T: StateData + StageLabel,
{
self.add_system_set(system_set.with_run_criteria(crate::in_state::<T>(state)));
self
}
}