diff --git a/crates/bevy_state/src/lib.rs b/crates/bevy_state/src/lib.rs index 88598db0b9ba7..c2aba9c24d1ac 100644 --- a/crates/bevy_state/src/lib.rs +++ b/crates/bevy_state/src/lib.rs @@ -44,6 +44,10 @@ pub mod state; /// Provides [`StateScoped`](crate::state_scoped::StateScoped) and /// [`clear_state_scoped_entities`](crate::state_scoped::clear_state_scoped_entities) for managing lifetime of entities. pub mod state_scoped; +#[cfg(feature = "bevy_app")] +/// Provides [`App`](bevy_app::App) and [`SubApp`](bevy_app::SubApp) with methods for registering +/// state-scoped events. +pub mod state_scoped_events; #[cfg(feature = "bevy_reflect")] /// Provides definitions for the basic traits required by the state system @@ -71,4 +75,7 @@ pub mod prelude { }; #[doc(hidden)] pub use crate::state_scoped::StateScoped; + #[cfg(feature = "bevy_app")] + #[doc(hidden)] + pub use crate::state_scoped_events::StateScopedEventsAppExt; } diff --git a/crates/bevy_state/src/state_scoped_events.rs b/crates/bevy_state/src/state_scoped_events.rs new file mode 100644 index 0000000000000..b3baeab3ca7f0 --- /dev/null +++ b/crates/bevy_state/src/state_scoped_events.rs @@ -0,0 +1,109 @@ +use std::marker::PhantomData; + +use bevy_app::{App, SubApp}; +use bevy_ecs::{ + event::{Event, EventReader, Events}, + system::{Commands, Resource}, + world::World, +}; +use bevy_utils::HashMap; + +use crate::state::{FreelyMutableState, OnExit, StateTransitionEvent}; + +fn clear_event_queue(w: &mut World) { + if let Some(mut queue) = w.get_resource_mut::>() { + queue.clear(); + } +} + +#[derive(Resource)] +struct StateScopedEvents { + cleanup_fns: HashMap>, +} + +impl StateScopedEvents { + fn add_event(&mut self, state: S) { + self.cleanup_fns + .entry(state) + .or_default() + .push(clear_event_queue::); + } + + fn cleanup(&self, w: &mut World, state: S) { + let Some(fns) = self.cleanup_fns.get(&state) else { + return; + }; + for callback in fns { + (*callback)(w); + } + } +} + +impl Default for StateScopedEvents { + fn default() -> Self { + Self { + cleanup_fns: HashMap::default(), + } + } +} + +fn cleanup_state_scoped_event( + mut c: Commands, + mut transitions: EventReader>, +) { + let Some(transition) = transitions.read().last() else { + return; + }; + if transition.entered == transition.exited { + return; + } + let Some(exited) = transition.exited.clone() else { + return; + }; + + c.add(move |w: &mut World| { + w.resource_scope::, ()>(|w, events| { + events.cleanup(w, exited); + }); + }); +} + +fn add_state_scoped_event_impl( + app: &mut SubApp, + _p: PhantomData, + state: S, +) { + if !app.world().contains_resource::>() { + app.init_resource::>(); + } + app.add_event::(); + app.world_mut() + .resource_mut::>() + .add_event::(state.clone()); + app.add_systems(OnExit(state), cleanup_state_scoped_event::); +} + +/// Extension trait for [`App`] adding methods for registering state scoped events. +pub trait StateScopedEventsAppExt { + /// Adds an [`Event`] that is automatically cleaned up when leaving the specified `state`. + /// + /// Note that event cleanup is ordered ambiguously relative to [`StateScoped`](crate::prelude::StateScoped) entity + /// cleanup and the [`OnExit`] schedule for the target state. All of these (state scoped + /// entities and events cleanup, and `OnExit`) occur within schedule [`StateTransition`](crate::prelude::StateTransition) + /// and system set `StateTransitionSteps::ExitSchedules`. + fn add_state_scoped_event(&mut self, state: impl FreelyMutableState) -> &mut Self; +} + +impl StateScopedEventsAppExt for App { + fn add_state_scoped_event(&mut self, state: impl FreelyMutableState) -> &mut Self { + add_state_scoped_event_impl(self.main_mut(), PhantomData::, state); + self + } +} + +impl StateScopedEventsAppExt for SubApp { + fn add_state_scoped_event(&mut self, state: impl FreelyMutableState) -> &mut Self { + add_state_scoped_event_impl(self, PhantomData::, state); + self + } +}