Refactor state scoped events to match entities.

crates/bevy_state/src/state_scoped_events.rs

@@ -10,29 +10,43 @@ use bevy_ecs::{
 };
 use bevy_platform::collections::HashMap;
 
-use crate::state::{OnExit, StateTransitionEvent, States};
+use crate::state::{OnEnter, OnExit, StateTransitionEvent, States};
 
 fn clear_event_queue<E: Event>(w: &mut World) {
     if let Some(mut queue) = w.get_resource_mut::<Events<E>>() {
         queue.clear();
     }
 }
 
+#[derive(Copy, Clone)]
+enum TransitionType {
+    OnExit,
+    OnEnter,
+}
+
 #[derive(Resource)]
 struct StateScopedEvents<S: States> {
-    cleanup_fns: HashMap<S, Vec<fn(&mut World)>>,
+    /// Keeps track of which events need to be reset when the state is exited.
+    on_exit: HashMap<S, Vec<fn(&mut World)>>,
+    /// Keeps track of which events need to be reset when the state is entered.
+    on_enter: HashMap<S, Vec<fn(&mut World)>>,
 }
 
 impl<S: States> StateScopedEvents<S> {
-    fn add_event<E: Event>(&mut self, state: S) {
-        self.cleanup_fns
-            .entry(state)
-            .or_default()
-            .push(clear_event_queue::<E>);
+    fn add_event<E: Event>(&mut self, state: S, transition_type: TransitionType) {
+        let map = match transition_type {
+            TransitionType::OnExit => &mut self.on_exit,
+            TransitionType::OnEnter => &mut self.on_enter,
+        };
+        map.entry(state).or_default().push(clear_event_queue::<E>);
     }
 
-    fn cleanup(&self, w: &mut World, state: S) {
-        let Some(fns) = self.cleanup_fns.get(&state) else {
+    fn cleanup(&self, w: &mut World, state: S, transition_type: TransitionType) {
+        let map = match transition_type {
+            TransitionType::OnExit => &self.on_exit,
+            TransitionType::OnEnter => &self.on_enter,
+        };
+        let Some(fns) = map.get(&state) else {
             return;
         };
         for callback in fns {
@@ -44,12 +58,13 @@ impl<S: States> StateScopedEvents<S> {
 impl<S: States> Default for StateScopedEvents<S> {
     fn default() -> Self {
         Self {
-            cleanup_fns: HashMap::default(),
+            on_exit: HashMap::default(),
+            on_enter: HashMap::default(),
         }
     }
 }
 
-fn cleanup_state_scoped_event<S: States>(
+fn clear_events_on_exit_state<S: States>(
     mut c: Commands,
     mut transitions: EventReader<StateTransitionEvent<S>>,
 ) {
@@ -65,48 +80,185 @@ fn cleanup_state_scoped_event<S: States>(
 
     c.queue(move |w: &mut World| {
         w.resource_scope::<StateScopedEvents<S>, ()>(|w, events| {
-            events.cleanup(w, exited);
+            events.cleanup(w, exited, TransitionType::OnExit);
         });
     });
 }
 
-fn add_state_scoped_event_impl<E: Event, S: States>(
+fn clear_events_on_enter_state<S: States>(
+    mut c: Commands,
+    mut transitions: EventReader<StateTransitionEvent<S>>,
+) {
+    let Some(transition) = transitions.read().last() else {
+        return;
+    };
+    if transition.entered == transition.exited {
+        return;
+    }
+    let Some(entered) = transition.entered.clone() else {
+        return;
+    };
+
+    c.queue(move |w: &mut World| {
+        w.resource_scope::<StateScopedEvents<S>, ()>(|w, events| {
+            events.cleanup(w, entered, TransitionType::OnEnter);
+        });
+    });
+}
+
+fn clear_events_on_state_transition<E: Event, S: States>(
     app: &mut SubApp,
     _p: PhantomData<E>,
     state: S,
+    transition_type: TransitionType,
 ) {
     if !app.world().contains_resource::<StateScopedEvents<S>>() {
         app.init_resource::<StateScopedEvents<S>>();
     }
-    app.add_event::<E>();
     app.world_mut()
         .resource_mut::<StateScopedEvents<S>>()
-        .add_event::<E>(state.clone());
-    app.add_systems(OnExit(state), cleanup_state_scoped_event::<S>);
+        .add_event::<E>(state.clone(), transition_type);
+    match transition_type {
+        TransitionType::OnExit => app.add_systems(OnExit(state), clear_events_on_exit_state::<S>),
+        TransitionType::OnEnter => {
+            app.add_systems(OnEnter(state), clear_events_on_enter_state::<S>)
+        }
+    };
 }
 
 /// 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`.
+    /// Clears an [`Event`] when exiting the specified `state`.
     ///
-    /// Note that event cleanup is ordered ambiguously relative to [`DespawnOnEnterState`](crate::prelude::DespawnOnEnterState)
-    /// and [`DespawnOnExitState`](crate::prelude::DespawnOnExitState) 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)
+    /// Note that event cleanup is ambiguously ordered relative to  
+    /// [`DespawnOnExitState`](crate::prelude::DespawnOnExitState) 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 `StateTransitionSystems::ExitSchedules`.
-    fn add_state_scoped_event<E: Event>(&mut self, state: impl States) -> &mut Self;
+    fn clear_events_on_exit_state<E: Event>(&mut self, state: impl States) -> &mut Self;
+
+    /// Clears an [`Event`] when entering the specified `state`.
+    ///
+    /// Note that event cleanup is ambiguously ordered relative to
+    /// [`DespawnOnEnterState`](crate::prelude::DespawnOnEnterState) entity cleanup,
+    /// and the [`OnEnter`] schedule for the target state.
+    /// All of these (state scoped entities and events cleanup, and `OnEnter`)
+    /// occur within schedule [`StateTransition`](crate::prelude::StateTransition)
+    /// and system set `StateTransitionSystems::EnterSchedules`.
+    fn clear_events_on_enter_state<E: Event>(&mut self, state: impl States) -> &mut Self;
 }
 
 impl StateScopedEventsAppExt for App {
-    fn add_state_scoped_event<E: Event>(&mut self, state: impl States) -> &mut Self {
-        add_state_scoped_event_impl(self.main_mut(), PhantomData::<E>, state);
+    fn clear_events_on_exit_state<E: Event>(&mut self, state: impl States) -> &mut Self {
+        clear_events_on_state_transition(
+            self.main_mut(),
+            PhantomData::<E>,
+            state,
+            TransitionType::OnExit,
+        );
+        self
+    }
+
+    fn clear_events_on_enter_state<E: Event>(&mut self, state: impl States) -> &mut Self {
+        clear_events_on_state_transition(
+            self.main_mut(),
+            PhantomData::<E>,
+            state,
+            TransitionType::OnEnter,
+        );
         self
     }
 }
 
 impl StateScopedEventsAppExt for SubApp {
-    fn add_state_scoped_event<E: Event>(&mut self, state: impl States) -> &mut Self {
-        add_state_scoped_event_impl(self, PhantomData::<E>, state);
+    fn clear_events_on_exit_state<E: Event>(&mut self, state: impl States) -> &mut Self {
+        clear_events_on_state_transition(self, PhantomData::<E>, state, TransitionType::OnExit);
+        self
+    }
+
+    fn clear_events_on_enter_state<E: Event>(&mut self, state: impl States) -> &mut Self {
+        clear_events_on_state_transition(self, PhantomData::<E>, state, TransitionType::OnEnter);
         self
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::app::StatesPlugin;
+    use bevy_state::prelude::*;
+
+    #[derive(States, Default, Clone, Hash, Eq, PartialEq, Debug)]
+    enum TestState {
+        #[default]
+        A,
+        B,
+    }
+
+    #[derive(Event, Debug)]
+    struct StandardEvent;
+
+    #[derive(Event, Debug)]
+    struct StateScopedEvent;
+
+    #[test]
+    fn clear_event_on_exit_state() {
+        let mut app = App::new();
+        app.add_plugins(StatesPlugin);
+        app.init_state::<TestState>();
+
+        app.add_event::<StandardEvent>();
+        app.add_event::<StateScopedEvent>()
+            .clear_events_on_exit_state::<StateScopedEvent>(TestState::A);
+
+        app.world_mut().send_event(StandardEvent).unwrap();
+        app.world_mut().send_event(StateScopedEvent).unwrap();
+        assert!(!app.world().resource::<Events<StandardEvent>>().is_empty());
+        assert!(!app
+            .world()
+            .resource::<Events<StateScopedEvent>>()
+            .is_empty());
+
+        app.world_mut()
+            .resource_mut::<NextState<TestState>>()
+            .set(TestState::B);
+        app.update();
+
+        assert!(!app.world().resource::<Events<StandardEvent>>().is_empty());
+        assert!(app
+            .world()
+            .resource::<Events<StateScopedEvent>>()
+            .is_empty());
+    }
+
+    #[test]
+    fn clear_event_on_enter_state() {
+        let mut app = App::new();
+        app.add_plugins(StatesPlugin);
+        app.init_state::<TestState>();
+
+        app.add_event::<StandardEvent>();
+        app.add_event::<StateScopedEvent>()
+            .clear_events_on_enter_state::<StateScopedEvent>(TestState::B);
+
+        app.world_mut().send_event(StandardEvent).unwrap();
+        app.world_mut().send_event(StateScopedEvent).unwrap();
+        assert!(!app.world().resource::<Events<StandardEvent>>().is_empty());
+        assert!(!app
+            .world()
+            .resource::<Events<StateScopedEvent>>()
+            .is_empty());
+
+        app.world_mut()
+            .resource_mut::<NextState<TestState>>()
+            .set(TestState::B);
+        app.update();
+
+        assert!(!app.world().resource::<Events<StandardEvent>>().is_empty());
+        assert!(app
+            .world()
+            .resource::<Events<StateScopedEvent>>()
+            .is_empty());
+    }
+}

release-content/migration-guides/rename_state_scoped.md

@@ -0,0 +1,20 @@
+---
+title: Renamed state scoped entities and events
+pull_requests: [18818, 19435]
+---
+
+Previously, Bevy provided the `StateScoped` component and `add_state_scoped_event` method
+as a way to remove entities/events when **exiting** a state.
+
+However, it can also be useful to have the opposite behavior,
+where entities/events are removed when **entering** a state.
+This is now possible with the new `DespawnOnEnterState` component and `clear_events_on_enter_state` method.
+
+To support this addition, the previous method and component have been renamed.
+Also, `clear_event_on_exit_state` no longer adds the event automatically, so you must call `App::add_event` manually.
+
+| Before                        | After                                      |
+|-------------------------------|--------------------------------------------|
+| `StateScoped`                 | `DespawnOnExitState`                       |
+| `clear_state_scoped_entities` | `despawn_entities_on_exit_state`           |
+| `add_state_scoped_event`      | `add_event` + `clear_events_on_exit_state` |