bevy_state_stack 0.2.1

An improved state stack for bevy.
Documentation
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
	}

	/// Add a system set to run when the top state is `state`
	/// Replaces the run criteria of the system set
	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
	}
}