Generalize StateTransitionEvent<S> to allow identity transitions (#13579)

# Objective

This PR addresses one of the issues from [discord state
discussion](https://discord.com/channels/691052431525675048/1237949214017716356).
Same-state transitions can be desirable, so there should exist a hook
for them.

Fixes https://github.com/bevyengine/bevy/issues/9130.

## Solution

- Allow `StateTransitionEvent<S>` to contain identity transitions.
- Ignore identity transitions at schedule running level (`OnExit`,
`OnTransition`, `OnEnter`).
- Propagate identity transitions through `SubStates` and
`ComputedStates`.
- Add example about registering custom transition schedules.

## Changelog

- `StateTransitionEvent<S>` can be emitted with same `exited` and
`entered` state.

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>

Cargo.toml

@@ -1771,14 +1771,14 @@ category = "ECS (Entity Component System)"
 wasm = false
 
 [[example]]
-name = "state"
+name = "states"
-path = "examples/state/state.rs"
+path = "examples/state/states.rs"
 doc-scrape-examples = true
 required-features = ["bevy_dev_tools"]
 
-[package.metadata.example.state]
+[package.metadata.example.states]
-name = "State"
+name = "States"
-description = "Illustrates how to use States to control transitioning from a Menu state to an InGame state"
+description = "Illustrates how to use States to control transitioning from a Menu state to an InGame state."
 category = "State"
 wasm = false
 
@@ -1802,7 +1802,19 @@ required-features = ["bevy_dev_tools"]
 
 [package.metadata.example.computed_states]
 name = "Computed States"
-description = "Advanced state patterns using Computed States"
+description = "Advanced state patterns using Computed States."
+category = "State"
+wasm = false
+
+[[example]]
+name = "custom_transitions"
+path = "examples/state/custom_transitions.rs"
+doc-scrape-examples = true
+required-features = ["bevy_dev_tools"]
+
+[package.metadata.example.custom_transitions]
+name = "Custom State Transition Behavior"
+description = "Creating and working with custom state transition schedules."
 category = "State"
 wasm = false
 

crates/bevy_state/src/lib.rs

@@ -47,8 +47,8 @@ pub mod prelude {
     pub use crate::condition::*;
     #[doc(hidden)]
     pub use crate::state::{
-        ComputedStates, NextState, OnEnter, OnExit, OnTransition, State, StateSet, StateTransition,
+        last_transition, ComputedStates, NextState, OnEnter, OnExit, OnTransition, State, StateSet,
-        StateTransitionEvent, States, SubStates,
+        StateTransition, StateTransitionEvent, StateTransitionSteps, States, SubStates,
     };
     #[doc(hidden)]
     pub use crate::state_scoped::StateScoped;

crates/bevy_state/src/state/freely_mutable_state.rs

@@ -54,7 +54,5 @@ fn apply_state_transition<S: FreelyMutableState>(
     let Some(current_state) = current_state else {
         return;
     };
-    if next_state != *current_state.get() {
     internal_apply_state_transition(event, commands, Some(current_state), Some(next_state));
-    }
 }

crates/bevy_state/src/state/mod.rs

@@ -497,4 +497,126 @@ mod tests {
             "Should Only Exit Twice"
         );
     }
+
+    #[derive(Resource, Default, PartialEq, Debug)]
+    struct TransitionCounter {
+        exit: u8,
+        transition: u8,
+        enter: u8,
+    }
+
+    #[test]
+    fn same_state_transition_should_emit_event_and_not_run_schedules() {
+        let mut world = World::new();
+        EventRegistry::register_event::<StateTransitionEvent<SimpleState>>(&mut world);
+        world.init_resource::<State<SimpleState>>();
+        let mut schedules = Schedules::new();
+        let mut apply_changes = Schedule::new(StateTransition);
+        SimpleState::register_state(&mut apply_changes);
+        schedules.insert(apply_changes);
+
+        world.insert_resource(TransitionCounter::default());
+        let mut on_exit = Schedule::new(OnExit(SimpleState::A));
+        on_exit.add_systems(|mut c: ResMut<TransitionCounter>| c.exit += 1);
+        schedules.insert(on_exit);
+        let mut on_transition = Schedule::new(OnTransition {
+            exited: SimpleState::A,
+            entered: SimpleState::A,
+        });
+        on_transition.add_systems(|mut c: ResMut<TransitionCounter>| c.transition += 1);
+        schedules.insert(on_transition);
+        let mut on_enter = Schedule::new(OnEnter(SimpleState::A));
+        on_enter.add_systems(|mut c: ResMut<TransitionCounter>| c.enter += 1);
+        schedules.insert(on_enter);
+
+        world.insert_resource(schedules);
+        setup_state_transitions_in_world(&mut world, None);
+
+        world.run_schedule(StateTransition);
+        assert_eq!(world.resource::<State<SimpleState>>().0, SimpleState::A);
+        assert!(world
+            .resource::<Events<StateTransitionEvent<SimpleState>>>()
+            .is_empty());
+
+        world.insert_resource(TransitionCounter::default());
+        world.insert_resource(NextState::Pending(SimpleState::A));
+        world.run_schedule(StateTransition);
+        assert_eq!(world.resource::<State<SimpleState>>().0, SimpleState::A);
+        assert_eq!(
+            *world.resource::<TransitionCounter>(),
+            TransitionCounter {
+                exit: 0,
+                transition: 0,
+                enter: 0
+            }
+        );
+        assert_eq!(
+            world
+                .resource::<Events<StateTransitionEvent<SimpleState>>>()
+                .len(),
+            1
+        );
+    }
+
+    #[test]
+    fn same_state_transition_should_propagate_to_sub_state() {
+        let mut world = World::new();
+        EventRegistry::register_event::<StateTransitionEvent<SimpleState>>(&mut world);
+        EventRegistry::register_event::<StateTransitionEvent<SubState>>(&mut world);
+        world.insert_resource(State(SimpleState::B(true)));
+        world.init_resource::<State<SubState>>();
+        let mut schedules = Schedules::new();
+        let mut apply_changes = Schedule::new(StateTransition);
+        SimpleState::register_state(&mut apply_changes);
+        SubState::register_sub_state_systems(&mut apply_changes);
+        schedules.insert(apply_changes);
+        world.insert_resource(schedules);
+        setup_state_transitions_in_world(&mut world, None);
+
+        world.insert_resource(NextState::Pending(SimpleState::B(true)));
+        world.run_schedule(StateTransition);
+        assert_eq!(
+            world
+                .resource::<Events<StateTransitionEvent<SimpleState>>>()
+                .len(),
+            1
+        );
+        assert_eq!(
+            world
+                .resource::<Events<StateTransitionEvent<SubState>>>()
+                .len(),
+            1
+        );
+    }
+
+    #[test]
+    fn same_state_transition_should_propagate_to_computed_state() {
+        let mut world = World::new();
+        EventRegistry::register_event::<StateTransitionEvent<SimpleState>>(&mut world);
+        EventRegistry::register_event::<StateTransitionEvent<TestComputedState>>(&mut world);
+        world.insert_resource(State(SimpleState::B(true)));
+        world.insert_resource(State(TestComputedState::BisTrue));
+        let mut schedules = Schedules::new();
+        let mut apply_changes = Schedule::new(StateTransition);
+        SimpleState::register_state(&mut apply_changes);
+        TestComputedState::register_computed_state_systems(&mut apply_changes);
+        schedules.insert(apply_changes);
+        world.insert_resource(schedules);
+        setup_state_transitions_in_world(&mut world, None);
+
+        world.insert_resource(NextState::Pending(SimpleState::B(true)));
+        world.run_schedule(StateTransition);
+        assert_eq!(
+            world
+                .resource::<Events<StateTransitionEvent<SimpleState>>>()
+                .len(),
+            1
+        );
+        assert_eq!(
+            world
+                .resource::<Events<StateTransitionEvent<TestComputedState>>>()
+                .len(),
+            1
+        );
+    }
 }

crates/bevy_state/src/state/transitions.rs

@@ -11,20 +11,24 @@ use bevy_ecs::{
 
 use super::{resources::State, states::States};
 
-/// The label of a [`Schedule`] that runs whenever [`State<S>`]
+/// The label of a [`Schedule`] that **only** runs whenever [`State<S>`] enters the provided state.
-/// enters this state.
+///
+/// This schedule ignores identity transitions.
 #[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
 pub struct OnEnter<S: States>(pub S);
 
-/// The label of a [`Schedule`] that runs whenever [`State<S>`]
+/// The label of a [`Schedule`] that **only** runs whenever [`State<S>`] exits the provided state.
-/// exits this state.
+///
+/// This schedule ignores identity transitions.
 #[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
 pub struct OnExit<S: States>(pub S);
 
 /// The label of a [`Schedule`] that **only** runs whenever [`State<S>`]
-/// exits the `from` state, AND enters the `to` state.
+/// exits AND enters the provided `exited` and `entered` states.
 ///
 /// Systems added to this schedule are always ran *after* [`OnExit`], and *before* [`OnEnter`].
+///
+/// This schedule will run on identity transitions.
 #[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
 pub struct OnTransition<S: States> {
     /// The state being exited.
@@ -38,6 +42,7 @@ pub struct OnTransition<S: States> {
 pub struct StateTransition;
 
 /// Event sent when any state transition of `S` happens.
+/// This includes identity transitions, where `exited` and `entered` have the same value.
 ///
 /// If you know exactly what state you want to respond to ahead of time, consider [`OnEnter`], [`OnTransition`], or [`OnExit`]
 #[derive(Debug, Copy, Clone, PartialEq, Eq, Event)]
@@ -52,11 +57,16 @@ pub struct StateTransitionEvent<S: States> {
 ///
 /// These system sets are run sequentially, in the order of the enum variants.
 #[derive(SystemSet, Clone, Debug, PartialEq, Eq, Hash)]
-pub(crate) enum StateTransitionSteps {
+pub enum StateTransitionSteps {
+    /// Parentless states apply their [`NextState<S>`].
     RootTransitions,
+    /// States with parents apply their computation and [`NextState<S>`].
     DependentTransitions,
+    /// Exit schedules are executed.
     ExitSchedules,
+    /// Transition schedules are executed.
     TransitionSchedules,
+    /// Enter schedules are executed.
     EnterSchedules,
 }
 
@@ -88,15 +98,18 @@ pub(crate) fn internal_apply_state_transition<S: States>(
                 // entering - we need to set the new value, compute dependant states, send transition events
                 // and register transition schedules.
                 Some(mut state_resource) => {
-                    if *state_resource != entered {
+                    let exited = match *state_resource == entered {
-                        let exited = mem::replace(&mut state_resource.0, entered.clone());
+                        true => entered.clone(),
+                        false => mem::replace(&mut state_resource.0, entered.clone()),
+                    };
 
+                    // Transition events are sent even for same state transitions
+                    // Although enter and exit schedules are not run by default.
                     event.send(StateTransitionEvent {
                         exited: Some(exited.clone()),
                         entered: Some(entered.clone()),
                     });
                 }
-                }
                 None => {
                     // If the [`State<S>`] resource does not exist, we create it, compute dependant states, send a transition event and register the `OnEnter` schedule.
                     commands.insert_resource(State(entered.clone()));
@@ -169,6 +182,9 @@ pub(crate) fn run_enter<S: States>(
     let Some(transition) = transition.0 else {
         return;
     };
+    if transition.entered == transition.exited {
+        return;
+    }
     let Some(entered) = transition.entered else {
         return;
     };
@@ -183,6 +199,9 @@ pub(crate) fn run_exit<S: States>(
     let Some(transition) = transition.0 else {
         return;
     };
+    if transition.entered == transition.exited {
+        return;
+    }
     let Some(exited) = transition.exited else {
         return;
     };

crates/bevy_state/src/state_scoped.rs

@@ -70,6 +70,9 @@ pub fn clear_state_scoped_entities<S: States>(
     let Some(transition) = transitions.read().last() else {
         return;
     };
+    if transition.entered == transition.exited {
+        return;
+    }
     let Some(exited) = &transition.exited else {
         return;
     };

examples/README.md

@@ -373,8 +373,9 @@ Example | Description
 
 Example | Description
 --- | ---
-[Computed States](../examples/state/computed_states.rs) | Advanced state patterns using Computed States
+[Computed States](../examples/state/computed_states.rs) | Advanced state patterns using Computed States.
-[State](../examples/state/state.rs) | Illustrates how to use States to control transitioning from a Menu state to an InGame state
+[Custom State Transition Behavior](../examples/state/custom_transitions.rs) | Creating and working with custom state transition schedules.
+[States](../examples/state/states.rs) | Illustrates how to use States to control transitioning from a Menu state to an InGame state.
 [Sub States](../examples/state/sub_states.rs) | Using Sub States for hierarchical state handling.
 
 ## Stress Tests

examples/state/custom_transitions.rs

@@ -0,0 +1,287 @@
+//! This example illustrates how to register custom state transition behavior.
+//!
+//! In this case we are trying to add `OnReenter` and `OnReexit`
+//! which will work much like `OnEnter` and `OnExit`,
+//! but additionally trigger if the state changed into itself.
+//!
+//! While identity transitions exist internally in [`StateTransitionEvent`]s,
+//! the default schedules intentionally ignore them, as this behavior is not commonly needed or expected.
+//!
+//! While this example displays identity transitions for a single state,
+//! identity transitions are propagated through the entire state graph,
+//! meaning any change to parent state will be propagated to [`ComputedStates`] and [`SubStates`].
+
+use std::marker::PhantomData;
+
+use bevy::{
+    dev_tools::states::*, ecs::schedule::ScheduleLabel, prelude::*,
+    state::state::StateTransitionSteps,
+};
+
+use custom_transitions::*;
+
+#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Hash, States)]
+enum AppState {
+    #[default]
+    Menu,
+    InGame,
+}
+
+fn main() {
+    App::new()
+        .add_plugins(DefaultPlugins)
+        .init_state::<AppState>()
+        .add_systems(Startup, setup)
+        .add_systems(OnEnter(AppState::Menu), setup_menu)
+        .add_systems(Update, menu.run_if(in_state(AppState::Menu)))
+        .add_systems(OnExit(AppState::Menu), cleanup_menu)
+        // We will restart the game progress every time we re-enter into it.
+        .add_plugins(IdentityTransitionsPlugin::<AppState>::default())
+        .add_systems(OnReenter(AppState::InGame), setup_game)
+        .add_systems(OnReexit(AppState::InGame), teardown_game)
+        // Doing it this way allows us to restart the game without any additional in-between states.
+        .add_systems(
+            Update,
+            ((movement, change_color, trigger_game_restart).run_if(in_state(AppState::InGame)),),
+        )
+        .add_systems(Update, log_transitions::<AppState>)
+        .run();
+}
+
+/// This module provides the custom `OnReenter` and `OnReexit` transitions for easy installation.
+mod custom_transitions {
+    use crate::*;
+
+    /// The plugin registers the transitions for one specific state.
+    /// If you use this for multiple states consider:
+    /// - installing the plugin multiple times,
+    /// - create an [`App`] extension method that inserts
+    ///   those transitions during state installation.
+    #[derive(Default)]
+    pub struct IdentityTransitionsPlugin<S: States>(PhantomData<S>);
+
+    impl<S: States> Plugin for IdentityTransitionsPlugin<S> {
+        fn build(&self, app: &mut App) {
+            app.add_systems(
+                StateTransition,
+                // The internals can generate at most one transition event of specific type per frame.
+                // We take the latest one and clear the queue.
+                last_transition::<S>
+                    // We insert the optional event into our schedule runner.
+                    .pipe(run_reenter::<S>)
+                    // State transitions are handled in three ordered steps, exposed as system sets.
+                    // We can add our systems to them, which will run the corresponding schedules when they're evaluated.
+                    // These are:
+                    // - [`StateTransitionSteps::ExitSchedules`]
+                    // - [`StateTransitionSteps::TransitionSchedules`]
+                    // - [`StateTransitionSteps::EnterSchedules`]
+                    .in_set(StateTransitionSteps::EnterSchedules),
+            )
+            .add_systems(
+                StateTransition,
+                last_transition::<S>
+                    .pipe(run_reexit::<S>)
+                    .in_set(StateTransitionSteps::ExitSchedules),
+            );
+        }
+    }
+
+    /// Custom schedule that will behave like [`OnEnter`], but run on identity transitions.
+    #[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
+    pub struct OnReenter<S: States>(pub S);
+
+    /// Schedule runner which checks conditions and if they're right
+    /// runs out custom schedule.
+    fn run_reenter<S: States>(transition: In<Option<StateTransitionEvent<S>>>, world: &mut World) {
+        // We return early if no transition event happened.
+        let Some(transition) = transition.0 else {
+            return;
+        };
+
+        // If we wanted to ignore identity transitions,
+        // we'd compare `exited` and `entered` here,
+        // and return if they were the same.
+
+        // We check if we actually entered a state.
+        // A [`None`] would indicate that the state was removed from the world.
+        // This only happens in the case of [`SubStates`] and [`ComputedStates`].
+        let Some(entered) = transition.entered else {
+            return;
+        };
+
+        // If all conditions are valid, we run our custom schedule.
+        let _ = world.try_run_schedule(OnReenter(entered));
+
+        // If you want to overwrite the default `OnEnter` behavior to act like re-enter,
+        // you can do so by running the `OnEnter` schedule here. Note that you don't want
+        // to run `OnEnter` when the default behavior does so.
+        // ```
+        // if transition.entered != transition.exited {
+        //     return;
+        // }
+        // let _ = world.try_run_schedule(OnReenter(entered));
+        // ```
+    }
+
+    /// Custom schedule that will behave like [`OnExit`], but run on identity transitions.
+    #[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
+    pub struct OnReexit<S: States>(pub S);
+
+    fn run_reexit<S: States>(transition: In<Option<StateTransitionEvent<S>>>, world: &mut World) {
+        let Some(transition) = transition.0 else {
+            return;
+        };
+        let Some(exited) = transition.exited else {
+            return;
+        };
+
+        let _ = world.try_run_schedule(OnReexit(exited));
+    }
+}
+
+fn menu(
+    mut next_state: ResMut<NextState<AppState>>,
+    mut interaction_query: Query<
+        (&Interaction, &mut UiImage),
+        (Changed<Interaction>, With<Button>),
+    >,
+) {
+    for (interaction, mut image) in &mut interaction_query {
+        let color = &mut image.color;
+        match *interaction {
+            Interaction::Pressed => {
+                *color = PRESSED_BUTTON;
+                next_state.set(AppState::InGame);
+            }
+            Interaction::Hovered => {
+                *color = HOVERED_BUTTON;
+            }
+            Interaction::None => {
+                *color = NORMAL_BUTTON;
+            }
+        }
+    }
+}
+
+fn cleanup_menu(mut commands: Commands, menu_data: Res<MenuData>) {
+    commands.entity(menu_data.button_entity).despawn_recursive();
+}
+
+const SPEED: f32 = 100.0;
+fn movement(
+    time: Res<Time>,
+    input: Res<ButtonInput<KeyCode>>,
+    mut query: Query<&mut Transform, With<Sprite>>,
+) {
+    for mut transform in &mut query {
+        let mut direction = Vec3::ZERO;
+        if input.pressed(KeyCode::ArrowLeft) {
+            direction.x -= 1.0;
+        }
+        if input.pressed(KeyCode::ArrowRight) {
+            direction.x += 1.0;
+        }
+        if input.pressed(KeyCode::ArrowUp) {
+            direction.y += 1.0;
+        }
+        if input.pressed(KeyCode::ArrowDown) {
+            direction.y -= 1.0;
+        }
+
+        if direction != Vec3::ZERO {
+            transform.translation += direction.normalize() * SPEED * time.delta_seconds();
+        }
+    }
+}
+
+fn change_color(time: Res<Time>, mut query: Query<&mut Sprite>) {
+    for mut sprite in &mut query {
+        let new_color = LinearRgba {
+            blue: (time.elapsed_seconds() * 0.5).sin() + 2.0,
+            ..LinearRgba::from(sprite.color)
+        };
+
+        sprite.color = new_color.into();
+    }
+}
+
+// We can restart the game by pressing "R".
+// This will trigger an [`AppState::InGame`] -> [`AppState::InGame`]
+// transition, which will run our custom schedules.
+fn trigger_game_restart(
+    input: Res<ButtonInput<KeyCode>>,
+    mut next_state: ResMut<NextState<AppState>>,
+) {
+    if input.just_pressed(KeyCode::KeyR) {
+        // Although we are already in this state setting it again will generate an identity transition.
+        // While default schedules ignore those kinds of transitions, our custom schedules will react to them.
+        next_state.set(AppState::InGame);
+    }
+}
+
+fn setup(mut commands: Commands) {
+    commands.spawn(Camera2dBundle::default());
+}
+
+fn setup_game(mut commands: Commands, asset_server: Res<AssetServer>) {
+    commands.spawn(SpriteBundle {
+        texture: asset_server.load("branding/icon.png"),
+        ..default()
+    });
+}
+
+fn teardown_game(mut commands: Commands, player: Query<Entity, With<Sprite>>) {
+    commands.entity(player.single()).despawn();
+}
+
+#[derive(Resource)]
+struct MenuData {
+    pub button_entity: Entity,
+}
+
+const NORMAL_BUTTON: Color = Color::srgb(0.15, 0.15, 0.15);
+const HOVERED_BUTTON: Color = Color::srgb(0.25, 0.25, 0.25);
+const PRESSED_BUTTON: Color = Color::srgb(0.35, 0.75, 0.35);
+
+fn setup_menu(mut commands: Commands) {
+    let button_entity = commands
+        .spawn(NodeBundle {
+            style: Style {
+                // center button
+                width: Val::Percent(100.),
+                height: Val::Percent(100.),
+                justify_content: JustifyContent::Center,
+                align_items: AlignItems::Center,
+                ..default()
+            },
+            ..default()
+        })
+        .with_children(|parent| {
+            parent
+                .spawn(ButtonBundle {
+                    style: Style {
+                        width: Val::Px(150.),
+                        height: Val::Px(65.),
+                        // horizontally center child text
+                        justify_content: JustifyContent::Center,
+                        // vertically center child text
+                        align_items: AlignItems::Center,
+                        ..default()
+                    },
+                    image: UiImage::default().with_color(NORMAL_BUTTON),
+                    ..default()
+                })
+                .with_children(|parent| {
+                    parent.spawn(TextBundle::from_section(
+                        "Play",
+                        TextStyle {
+                            font_size: 40.0,
+                            color: Color::srgb(0.9, 0.9, 0.9),
+                            ..default()
+                        },
+                    ));
+                });
+        })
+        .id();
+    commands.insert_resource(MenuData { button_entity });
+}