Move utilities from examples to bevy_state and add concept of state-scoped entities (#13649)

# Objective

Move `StateScoped` and `log_transitions` to `bevy_state`, since they're
useful for end users.

Addresses #12852, although not in the way the issue had in mind.

## Solution

- Added `bevy_hierarchy` to default features of `bevy_state`.
- Move `log_transitions` to `transitions` module.
- Move `StateScoped` to `state_scoped` module, gated behind
`bevy_hierarchy` feature.
- Refreshed implementation.
- Added `enable_state_coped_entities<S: States>()` to add required
machinery to `App` for clearing state-scoped entities.


## Changelog

- Added `log_transitions` for displaying state transitions.
- Added `StateScoped` for binding entity lifetime to state and app
`enable_state_coped_entities` to register cleaning behavior.

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
Co-authored-by: François Mockers <francois.mockers@vleue.com>

Cargo.toml

@@ -1774,6 +1774,7 @@ wasm = false
 name = "state"
 path = "examples/state/state.rs"
 doc-scrape-examples = true
+required-features = ["bevy_dev_tools"]
 
 [package.metadata.example.state]
 name = "State"
@@ -1785,6 +1786,7 @@ wasm = false
 name = "sub_states"
 path = "examples/state/sub_states.rs"
 doc-scrape-examples = true
+required-features = ["bevy_dev_tools"]
 
 [package.metadata.example.sub_states]
 name = "Sub States"
@@ -1796,6 +1798,7 @@ wasm = false
 name = "computed_states"
 path = "examples/state/computed_states.rs"
 doc-scrape-examples = true
+required-features = ["bevy_dev_tools"]
 
 [package.metadata.example.computed_states]
 name = "Computed States"

crates/bevy_dev_tools/Cargo.toml

@@ -36,6 +36,7 @@ bevy_ui = { path = "../bevy_ui", version = "0.14.0-dev", features = [
 bevy_utils = { path = "../bevy_utils", version = "0.14.0-dev" }
 bevy_window = { path = "../bevy_window", version = "0.14.0-dev" }
 bevy_text = { path = "../bevy_text", version = "0.14.0-dev" }
+bevy_state = { path = "../bevy_state", version = "0.14.0-dev" }
 
 # other
 serde = { version = "1.0", features = ["derive"], optional = true }

crates/bevy_dev_tools/src/lib.rs

@@ -18,6 +18,8 @@ pub mod fps_overlay;
 #[cfg(feature = "bevy_ui_debug")]
 pub mod ui_debug_overlay;
 
+pub mod states;
+
 /// Enables developer tools in an [`App`]. This plugin is added automatically with `bevy_dev_tools`
 /// feature.
 ///

crates/bevy_dev_tools/src/states.rs

@@ -0,0 +1,18 @@
+//! Tools for debugging states.
+
+use bevy_ecs::event::EventReader;
+use bevy_state::state::{StateTransitionEvent, States};
+use bevy_utils::tracing::info;
+
+/// Logs state transitions into console.
+///
+/// This system is provided to make debugging easier by tracking state changes.
+pub fn log_transitions<S: States>(mut transitions: EventReader<StateTransitionEvent<S>>) {
+    // State internals can generate at most one event (of type) per frame.
+    let Some(transition) = transitions.read().last() else {
+        return;
+    };
+    let name = std::any::type_name::<S>();
+    let StateTransitionEvent { exited, entered } = transition;
+    info!("{} transition: {:?} => {:?}", name, exited, entered);
+}

crates/bevy_state/Cargo.toml

@@ -12,9 +12,10 @@ categories = ["game-engines", "data-structures"]
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [features]
-default = ["bevy_reflect", "bevy_app"]
+default = ["bevy_reflect", "bevy_app", "bevy_hierarchy"]
 bevy_reflect = ["dep:bevy_reflect", "bevy_ecs/bevy_reflect"]
 bevy_app = ["dep:bevy_app"]
+bevy_hierarchy = ["dep:bevy_hierarchy"]
 
 [dependencies]
 bevy_ecs = { path = "../bevy_ecs", version = "0.14.0-dev" }
@@ -22,6 +23,7 @@ bevy_state_macros = { path = "macros", version = "0.14.0-dev" }
 bevy_utils = { path = "../bevy_utils", version = "0.14.0-dev" }
 bevy_reflect = { path = "../bevy_reflect", version = "0.14.0-dev", optional = true }
 bevy_app = { path = "../bevy_app", version = "0.14.0-dev", optional = true }
+bevy_hierarchy = { path = "../bevy_hierarchy", version = "0.14.0-dev", optional = true }
 
 [lints]
 workspace = true

crates/bevy_state/src/app.rs

@@ -1,10 +1,15 @@
 use bevy_app::{App, MainScheduleOrder, Plugin, PreUpdate, Startup, SubApp};
-use bevy_ecs::{event::Events, schedule::ScheduleLabel, world::FromWorld};
+use bevy_ecs::{
+    event::Events,
+    schedule::{IntoSystemConfigs, ScheduleLabel},
+    world::FromWorld,
+};
 
 use crate::state::{
     setup_state_transitions_in_world, ComputedStates, FreelyMutableState, NextState, State,
-    StateTransition, StateTransitionEvent, SubStates,
+    StateTransition, StateTransitionEvent, StateTransitionSteps, States, SubStates,
 };
+use crate::state_scoped::clear_state_scoped_entities;
 
 /// State installation methods for [`App`](bevy_app::App) and [`SubApp`](bevy_app::SubApp).
 pub trait AppExtStates {
@@ -44,6 +49,11 @@ pub trait AppExtStates {
     ///
     /// This method is idempotent: it has no effect when called again using the same generic type.
     fn add_sub_state<S: SubStates>(&mut self) -> &mut Self;
+
+    /// Enable state-scoped entity clearing for state `S`.
+    ///
+    /// For more information refer to [`StateScoped`](crate::state_scoped::StateScoped).
+    fn enable_state_scoped_entities<S: States>(&mut self) -> &mut Self;
 }
 
 impl AppExtStates for SubApp {
@@ -91,10 +101,13 @@ impl AppExtStates for SubApp {
             self.add_event::<StateTransitionEvent<S>>();
             let schedule = self.get_schedule_mut(StateTransition).unwrap();
             S::register_computed_state_systems(schedule);
-            let state = self.world().resource::<State<S>>().get().clone();
+            let state = self
+                .world()
+                .get_resource::<State<S>>()
+                .map(|s| s.get().clone());
             self.world_mut().send_event(StateTransitionEvent {
                 exited: None,
-                entered: Some(state),
+                entered: state,
             });
         }
 
@@ -111,15 +124,36 @@ impl AppExtStates for SubApp {
             self.add_event::<StateTransitionEvent<S>>();
             let schedule = self.get_schedule_mut(StateTransition).unwrap();
             S::register_sub_state_systems(schedule);
-            let state = self.world().resource::<State<S>>().get().clone();
+            let state = self
+                .world()
+                .get_resource::<State<S>>()
+                .map(|s| s.get().clone());
             self.world_mut().send_event(StateTransitionEvent {
                 exited: None,
-                entered: Some(state),
+                entered: state,
             });
         }
 
         self
     }
+
+    fn enable_state_scoped_entities<S: States>(&mut self) -> &mut Self {
+        use bevy_utils::tracing::warn;
+
+        if !self
+            .world()
+            .contains_resource::<Events<StateTransitionEvent<S>>>()
+        {
+            let name = std::any::type_name::<S>();
+            warn!("State scoped entities are enabled for state `{}`, but the state isn't installed in the app!", name);
+        }
+        // We work with [`StateTransition`] in set [`StateTransitionSteps::ExitSchedules`] as opposed to [`OnExit`],
+        // because [`OnExit`] only runs for one specific variant of the state.
+        self.add_systems(
+            StateTransition,
+            clear_state_scoped_entities::<S>.in_set(StateTransitionSteps::ExitSchedules),
+        )
+    }
 }
 
 impl AppExtStates for App {
@@ -142,6 +176,12 @@ impl AppExtStates for App {
         self.main_mut().add_sub_state::<S>();
         self
     }
+
+    #[cfg(feature = "bevy_hierarchy")]
+    fn enable_state_scoped_entities<S: States>(&mut self) -> &mut Self {
+        self.main_mut().enable_state_scoped_entities::<S>();
+        self
+    }
 }
 
 /// Registers the [`StateTransition`] schedule in the [`MainScheduleOrder`] to enable state processing.

crates/bevy_state/src/lib.rs

@@ -35,6 +35,9 @@ pub mod condition;
 /// Provides definitions for the basic traits required by the state system
 pub mod state;
 
+/// Provides [`StateScoped`] and [`clear_state_scoped_entities`] for managing lifetime of entities.
+pub mod state_scoped;
+
 /// Most commonly used re-exported types.
 pub mod prelude {
     #[cfg(feature = "bevy_app")]
@@ -47,4 +50,6 @@ pub mod prelude {
         ComputedStates, NextState, OnEnter, OnExit, OnTransition, State, StateSet, StateTransition,
         StateTransitionEvent, States, SubStates,
     };
+    #[doc(hidden)]
+    pub use crate::state_scoped::StateScoped;
 }

crates/bevy_state/src/state_scoped.rs

@@ -0,0 +1,84 @@
+use bevy_ecs::{
+    component::Component,
+    entity::Entity,
+    event::EventReader,
+    system::{Commands, Query},
+};
+#[cfg(feature = "bevy_hierarchy")]
+use bevy_hierarchy::DespawnRecursiveExt;
+
+use crate::state::{StateTransitionEvent, States};
+
+/// Entities marked with this component will be removed
+/// when the world's state of the matching type no longer matches the supplied value.
+///
+/// To enable this feature remember to configure your application
+/// with [`enable_state_scoped_entities`](crate::app::AppExtStates::enable_state_scoped_entities) on your state(s) of choice.
+///
+/// If `bevy_hierarchy` feature is enabled, which it is by default, the despawn will be recursive.
+///
+/// ```
+/// use bevy_state::prelude::*;
+/// use bevy_ecs::prelude::*;
+///
+/// #[derive(Clone, Copy, PartialEq, Eq, Hash, Debug, Default, States)]
+/// enum GameState {
+///     #[default]
+///     MainMenu,
+///     SettingsMenu,
+///     InGame,
+/// }
+///
+/// # #[derive(Component)]
+/// # struct Player;
+///
+/// fn spawn_player(mut commands: Commands) {
+///     commands.spawn((
+///         StateScoped(GameState::InGame),
+///         Player
+///     ));
+/// }
+///
+/// # struct AppMock;
+/// # impl AppMock {
+/// #     fn init_state<S>(&mut self) {}
+/// #     fn enable_state_scoped_entities<S>(&mut self) {}
+/// #     fn add_systems<S, M>(&mut self, schedule: S, systems: impl IntoSystemConfigs<M>) {}
+/// # }
+/// # struct Update;
+/// # let mut app = AppMock;
+///
+/// app.init_state::<GameState>();
+/// app.enable_state_scoped_entities::<GameState>();
+/// app.add_systems(OnEnter(GameState::InGame), spawn_player);
+/// ```
+#[derive(Component)]
+pub struct StateScoped<S: States>(pub S);
+
+/// Removes entities marked with [`StateScoped<S>`]
+/// when their state no longer matches the world state.
+///
+/// If `bevy_hierarchy` feature is enabled, which it is by default, the despawn will be recursive.
+pub fn clear_state_scoped_entities<S: States>(
+    mut commands: Commands,
+    mut transitions: EventReader<StateTransitionEvent<S>>,
+    query: Query<(Entity, &StateScoped<S>)>,
+) {
+    // We use the latest event, because state machine internals generate at most 1
+    // transition event (per type) each frame. No event means no change happened
+    // and we skip iterating all entities.
+    let Some(transition) = transitions.read().last() else {
+        return;
+    };
+    let Some(exited) = &transition.exited else {
+        return;
+    };
+    for (entity, binding) in &query {
+        if binding.0 == *exited {
+            #[cfg(feature = "bevy_hierarchy")]
+            commands.entity(entity).despawn_recursive();
+            #[cfg(not(feature = "bevy_hierarchy"))]
+            commands.entity(entity).despawn();
+        }
+    }
+}

examples/state/computed_states.rs

@@ -16,7 +16,7 @@
 //! And lastly, we'll add [`Tutorial`], a computed state deriving from [`TutorialState`], [`InGame`] and [`IsPaused`], with 2 distinct
 //! states to display the 2 tutorial texts.
 
-use bevy::prelude::*;
+use bevy::{dev_tools::states::*, prelude::*};
 
 use ui::*;
 
@@ -186,7 +186,7 @@ fn main() {
         .add_systems(OnEnter(InGame), setup_game)
         // And we only want to run the [`clear_game`] function when we leave the [`AppState::InGame`] state, regardless
         // of whether we're paused.
-        .add_systems(OnExit(InGame), clear_state_bound_entities(InGame))
+        .enable_state_scoped_entities::<InGame>()
         // We want the color change, toggle_pause and quit_to_menu systems to ignore the paused condition, so we can use the [`InGame`] derived
         // state here as well.
         .add_systems(
@@ -200,26 +200,22 @@ fn main() {
         )
         // We can continue setting things up, following all the same patterns used above and in the `states` example.
         .add_systems(OnEnter(IsPaused::Paused), setup_paused_screen)
-        .add_systems(
+        .enable_state_scoped_entities::<IsPaused>()
-            OnExit(IsPaused::Paused),
-            clear_state_bound_entities(IsPaused::Paused),
-        )
         .add_systems(OnEnter(TurboMode), setup_turbo_text)
-        .add_systems(OnExit(TurboMode), clear_state_bound_entities(TurboMode))
+        .enable_state_scoped_entities::<TurboMode>()
         .add_systems(
             OnEnter(Tutorial::MovementInstructions),
             movement_instructions,
         )
         .add_systems(OnEnter(Tutorial::PauseInstructions), pause_instructions)
+        .enable_state_scoped_entities::<Tutorial>()
         .add_systems(
-            OnExit(Tutorial::MovementInstructions),
+            Update,
-            clear_state_bound_entities(Tutorial::MovementInstructions),
+            (
+                log_transitions::<AppState>,
+                log_transitions::<TutorialState>,
+            ),
         )
-        .add_systems(
-            OnExit(Tutorial::PauseInstructions),
-            clear_state_bound_entities(Tutorial::PauseInstructions),
-        )
-        .add_systems(Update, log_transitions)
         .run();
 }
 
@@ -277,22 +273,6 @@ fn menu(
     }
 }
 
-#[derive(Component)]
-struct StateBound<S: States>(S);
-
-fn clear_state_bound_entities<S: States>(
-    state: S,
-) -> impl Fn(Commands, Query<(Entity, &StateBound<S>)>) {
-    info!("Clearing entities for {state:?}");
-    move |mut commands, query| {
-        for (entity, bound) in &query {
-            if bound.0 == state {
-                commands.entity(entity).despawn_recursive();
-            }
-        }
-    }
-}
-
 fn toggle_pause(
     input: Res<ButtonInput<KeyCode>>,
     current_state: Res<State<AppState>>,
@@ -329,25 +309,6 @@ fn quit_to_menu(input: Res<ButtonInput<KeyCode>>, mut next_state: ResMut<NextSta
     }
 }
 
-/// print when either an `AppState` transition or a `TutorialState` transition happens
-fn log_transitions(
-    mut transitions: EventReader<StateTransitionEvent<AppState>>,
-    mut tutorial_transitions: EventReader<StateTransitionEvent<TutorialState>>,
-) {
-    for transition in transitions.read() {
-        info!(
-            "transition: {:?} => {:?}",
-            transition.exited, transition.entered
-        );
-    }
-    for transition in tutorial_transitions.read() {
-        info!(
-            "tutorial transition: {:?} => {:?}",
-            transition.exited, transition.entered
-        );
-    }
-}
-
 mod ui {
     use crate::*;
 
@@ -461,7 +422,7 @@ mod ui {
 
     pub fn setup_game(mut commands: Commands, asset_server: Res<AssetServer>) {
         commands.spawn((
-            StateBound(InGame),
+            StateScoped(InGame),
             SpriteBundle {
                 texture: asset_server.load("branding/icon.png"),
                 ..default()
@@ -505,7 +466,7 @@ mod ui {
         info!("Printing Pause");
         commands
             .spawn((
-                StateBound(IsPaused::Paused),
+                StateScoped(IsPaused::Paused),
                 NodeBundle {
                     style: Style {
                         // center button
@@ -555,7 +516,7 @@ mod ui {
     pub fn setup_turbo_text(mut commands: Commands) {
         commands
             .spawn((
-                StateBound(TurboMode),
+                StateScoped(TurboMode),
                 NodeBundle {
                     style: Style {
                         // center button
@@ -597,7 +558,7 @@ mod ui {
     pub fn movement_instructions(mut commands: Commands) {
         commands
             .spawn((
-                StateBound(Tutorial::MovementInstructions),
+                StateScoped(Tutorial::MovementInstructions),
                 NodeBundle {
                     style: Style {
                         // center button
@@ -654,7 +615,7 @@ mod ui {
     pub fn pause_instructions(mut commands: Commands) {
         commands
             .spawn((
-                StateBound(Tutorial::PauseInstructions),
+                StateScoped(Tutorial::PauseInstructions),
                 NodeBundle {
                     style: Style {
                         // center button

examples/state/state.rs

@@ -5,7 +5,7 @@
 //!
 //! In this case, we're transitioning from a `Menu` state to an `InGame` state.
 
-use bevy::prelude::*;
+use bevy::{dev_tools::states::*, prelude::*};
 
 fn main() {
     App::new()
@@ -25,7 +25,7 @@ fn main() {
             Update,
             (movement, change_color).run_if(in_state(AppState::InGame)),
         )
-        .add_systems(Update, log_transitions)
+        .add_systems(Update, log_transitions::<AppState>)
         .run();
 }
 
@@ -163,14 +163,3 @@ fn change_color(time: Res<Time>, mut query: Query<&mut Sprite>) {
         sprite.color = new_color.into();
     }
 }
-
-/// print when an `AppState` transition happens
-/// also serves as an example of how to use `StateTransitionEvent`
-fn log_transitions(mut transitions: EventReader<StateTransitionEvent<AppState>>) {
-    for transition in transitions.read() {
-        info!(
-            "transition: {:?} => {:?}",
-            transition.exited, transition.entered
-        );
-    }
-}

examples/state/sub_states.rs

@@ -7,7 +7,7 @@
 //! In this case, we're transitioning from a `Menu` state to an `InGame` state, at which point we create
 //! a substate called `IsPaused` to track whether the game is paused or not.
 
-use bevy::prelude::*;
+use bevy::{dev_tools::states::*, prelude::*};
 
 use ui::*;
 
@@ -43,10 +43,7 @@ fn main() {
         .add_systems(OnExit(AppState::Menu), cleanup_menu)
         .add_systems(OnEnter(AppState::InGame), setup_game)
         .add_systems(OnEnter(IsPaused::Paused), setup_paused_screen)
-        .add_systems(
+        .enable_state_scoped_entities::<IsPaused>()
-            OnExit(IsPaused::Paused),
-            clear_state_bound_entities(IsPaused::Paused),
-        )
         .add_systems(
             Update,
             (
@@ -59,7 +56,7 @@ fn main() {
                 toggle_pause.run_if(in_state(AppState::InGame)),
             ),
         )
-        .add_systems(Update, log_transitions)
+        .add_systems(Update, log_transitions::<AppState>)
         .run();
 }
 
@@ -142,31 +139,6 @@ fn toggle_pause(
     }
 }
 
-#[derive(Component)]
-struct StateBound<S: States>(S);
-
-fn clear_state_bound_entities<S: States>(
-    state: S,
-) -> impl Fn(Commands, Query<(Entity, &StateBound<S>)>) {
-    move |mut commands, query| {
-        for (entity, bound) in &query {
-            if bound.0 == state {
-                commands.entity(entity).despawn_recursive();
-            }
-        }
-    }
-}
-
-/// print when an `AppState` transition happens
-fn log_transitions(mut transitions: EventReader<StateTransitionEvent<AppState>>) {
-    for transition in transitions.read() {
-        info!(
-            "transition: {:?} => {:?}",
-            transition.exited, transition.entered
-        );
-    }
-}
-
 mod ui {
     use crate::*;
 
@@ -236,7 +208,7 @@ mod ui {
     pub fn setup_paused_screen(mut commands: Commands) {
         commands
             .spawn((
-                StateBound(IsPaused::Paused),
+                StateScoped(IsPaused::Paused),
                 NodeBundle {
                     style: Style {
                         // center button