diff --git a/crates/bevy_ecs/src/schedule/state.rs b/crates/bevy_ecs/src/schedule/state.rs index 1fa6f0b2d9450..f54a5f0ebaf1c 100644 --- a/crates/bevy_ecs/src/schedule/state.rs +++ b/crates/bevy_ecs/src/schedule/state.rs @@ -33,6 +33,21 @@ impl Default for StateStage { #[allow(clippy::mem_discriminant_non_enum)] impl StateStage { + pub fn with_enter_stage(mut self, state: T, stage: S) -> Self { + self.set_enter_stage(state, stage); + self + } + + pub fn with_exit_stage(mut self, state: T, stage: S) -> Self { + self.set_exit_stage(state, stage); + self + } + + pub fn with_update_stage(mut self, state: T, stage: S) -> Self { + self.set_update_stage(state, stage); + self + } + pub fn set_enter_stage(&mut self, state: T, stage: S) -> &mut Self { let stages = self.state_stages(state); stages.enter = Box::new(stage);