Add RefreshState resource and support same-state transitions

crates/bevy_app/src/sub_app.rs

@@ -337,6 +337,7 @@ impl SubApp {
             .contains_resource::<Events<StateTransitionEvent<S>>>()
         {
             setup_state_transitions_in_world(&mut self.world, Some(Startup.intern()));
+            self.init_resource::<RefreshState<S>>();
             self.add_event::<StateTransitionEvent<S>>();
             let schedule = self.get_schedule_mut(StateTransition).unwrap();
             S::register_computed_state_systems(schedule);

crates/bevy_state/src/lib.rs

@@ -38,7 +38,7 @@ pub mod prelude {
     pub use crate::condition::*;
     #[doc(hidden)]
     pub use crate::state::{
-        apply_state_transition, ComputedStates, NextState, OnEnter, OnExit, OnTransition, State,
-        StateSet, StateTransition, StateTransitionEvent, States, SubStates,
+        apply_state_transition, ComputedStates, NextState, OnEnter, OnExit, OnTransition,
+        RefreshState, State, StateSet, StateTransition, StateTransitionEvent, States, SubStates,
     };
 }

crates/bevy_state/src/state/resources.rs

@@ -1,11 +1,13 @@
-use std::ops::Deref;
+use std::{marker::PhantomData, ops::Deref};
 
 use bevy_ecs::{
     system::Resource,
     world::{FromWorld, World},
 };
 
-use super::{freely_mutable_state::FreelyMutableState, states::States};
+use super::{
+    computed_states::ComputedStates, freely_mutable_state::FreelyMutableState, states::States,
+};
 
 #[cfg(feature = "bevy_reflect")]
 use bevy_ecs::prelude::ReflectResource;
@@ -131,3 +133,28 @@ impl<S: FreelyMutableState> NextState<S> {
         *self = Self::Unchanged;
     }
 }
+
+/// The flag that determines whether the computed state `S` will be refreshed this frame.
+///
+/// Refreshing a state will apply its state transition even if nothing has changed, in which case
+/// there will be a state transition from the current state to itself.
+#[derive(Resource, Debug)]
+#[cfg_attr(
+    feature = "bevy_reflect",
+    derive(bevy_reflect::Reflect),
+    reflect(Resource)
+)]
+pub struct RefreshState<S: ComputedStates>(pub bool, PhantomData<S>);
+
+impl<S: ComputedStates> Default for RefreshState<S> {
+    fn default() -> Self {
+        Self(false, PhantomData)
+    }
+}
+
+impl<S: ComputedStates> RefreshState<S> {
+    /// Plan to refresh the computed state `S`.
+    pub fn refresh(&mut self) {
+        self.0 = true;
+    }
+}

crates/bevy_state/src/state/state_set.rs

@@ -10,7 +10,7 @@ use self::sealed::StateSetSealed;
 use super::{
     apply_state_transition, computed_states::ComputedStates, internal_apply_state_transition,
     run_enter, run_exit, run_transition, should_run_transition, sub_states::SubStates,
-    ApplyStateTransition, OnEnter, OnExit, OnTransition, State, StateTransitionEvent,
+    ApplyStateTransition, OnEnter, OnExit, OnTransition, RefreshState, State, StateTransitionEvent,
     StateTransitionSteps, States,
 };
 
@@ -57,7 +57,7 @@ pub trait StateSet: sealed::StateSetSealed {
 /// The isolation works because it is implemented for both S & [`Option<S>`], and has the `RawState` associated type
 /// that allows it to know what the resource in the world should be. We can then essentially "unwrap" it in our
 /// `StateSet` implementation - and the behaviour of that unwrapping will depend on the arguments expected by the
-/// the [`ComputedStates`] & [`SubStates]`.
+/// the [`ComputedStates`] & [`SubStates`].
 trait InnerStateSet: Sized {
     type RawState: States;
 
@@ -98,7 +98,9 @@ impl<S: InnerStateSet> StateSet for S {
                       event: EventWriter<StateTransitionEvent<T>>,
                       commands: Commands,
                       current_state: Option<ResMut<State<T>>>,
+                      refresh: Option<ResMut<RefreshState<T>>>,
                       state_set: Option<Res<State<S::RawState>>>| {
+            let should_refresh = refresh.is_some_and(|mut x| std::mem::take(&mut x.0));
             if parent_changed.is_empty() {
                 return;
             }
@@ -111,6 +113,11 @@ impl<S: InnerStateSet> StateSet for S {
                     None
                 };
 
+            let same_state = current_state.as_ref().map(|x| x.get()) == new_state.as_ref();
+            if same_state && !should_refresh {
+                return;
+            }
+
             internal_apply_state_transition(event, commands, current_state, new_state);
         };
 
@@ -158,21 +165,9 @@ impl<S: InnerStateSet> StateSet for S {
                     None
                 };
 
-            match new_state {
-                Some(value) => {
-                    if current_state.is_none() {
-                        internal_apply_state_transition(
-                            event,
-                            commands,
-                            current_state,
-                            Some(value),
-                        );
-                    }
-                }
-                None => {
-                    internal_apply_state_transition(event, commands, current_state, None);
-                }
-            };
+            if current_state.is_none() || new_state.is_none() {
+                internal_apply_state_transition(event, commands, current_state, new_state);
+            }
         };
 
         schedule
@@ -215,7 +210,13 @@ macro_rules! impl_state_set_sealed_tuples {
             fn register_computed_state_systems_in_schedule<T: ComputedStates<SourceStates = Self>>(
                 schedule: &mut Schedule,
             ) {
-                let system = |($(mut $evt),*,): ($(EventReader<StateTransitionEvent<$param::RawState>>),*,), event: EventWriter<StateTransitionEvent<T>>, commands: Commands, current_state: Option<ResMut<State<T>>>, ($($val),*,): ($(Option<Res<State<$param::RawState>>>),*,)| {
+                let system = |($(mut $evt),*,): ($(EventReader<StateTransitionEvent<$param::RawState>>),*,),
+                              event: EventWriter<StateTransitionEvent<T>>,
+                              commands: Commands,
+                              current_state: Option<ResMut<State<T>>>,
+                              refresh: Option<ResMut<RefreshState<T>>>,
+                              ($($val),*,): ($(Option<Res<State<$param::RawState>>>),*,)| {
+                    let should_refresh = refresh.is_some_and(|mut x| std::mem::take(&mut x.0));
                     if ($($evt.is_empty())&&*) {
                         return;
                     }
@@ -227,6 +228,11 @@ macro_rules! impl_state_set_sealed_tuples {
                         None
                     };
 
+                    let same_state = current_state.as_ref().map(|x| x.get()) == new_state.as_ref();
+                    if same_state && !should_refresh {
+                        return;
+                    }
+
                     internal_apply_state_transition(event, commands, current_state, new_state);
                 };
 
@@ -245,7 +251,11 @@ macro_rules! impl_state_set_sealed_tuples {
             fn register_sub_state_systems_in_schedule<T: SubStates<SourceStates = Self>>(
                 schedule: &mut Schedule,
             ) {
-                let system = |($(mut $evt),*,): ($(EventReader<StateTransitionEvent<$param::RawState>>),*,), event: EventWriter<StateTransitionEvent<T>>, commands: Commands, current_state: Option<ResMut<State<T>>>, ($($val),*,): ($(Option<Res<State<$param::RawState>>>),*,)| {
+                let system = |($(mut $evt),*,): ($(EventReader<StateTransitionEvent<$param::RawState>>),*,),
+                              event: EventWriter<StateTransitionEvent<T>>,
+                              commands: Commands,
+                              current_state: Option<ResMut<State<T>>>,
+                              ($($val),*,): ($(Option<Res<State<$param::RawState>>>),*,)| {
                     if ($($evt.is_empty())&&*) {
                         return;
                     }
@@ -256,16 +266,10 @@ macro_rules! impl_state_set_sealed_tuples {
                     } else {
                         None
                     };
-                    match new_state {
-                        Some(value) => {
-                            if current_state.is_none() {
-                                internal_apply_state_transition(event, commands, current_state, Some(value));
-                            }
-                        }
-                        None => {
-                            internal_apply_state_transition(event, commands, current_state, None);
-                        },
-                    };
+
+                    if current_state.is_none() || new_state.is_none() {
+                        internal_apply_state_transition(event, commands, current_state, new_state);
+                    }
                 };
 
                 schedule

crates/bevy_state/src/state/transitions.rs

@@ -74,8 +74,7 @@ impl<S: States> ApplyStateTransition<S> {
     }
 }
 
-/// This function actually applies a state change, and registers the required
-/// schedules for downstream computed states and transition schedules.
+/// This function updates the current state and sends a transition event.
 ///
 /// The `new_state` is an option to allow for removal - `None` will trigger the
 /// removal of the `State<S>` resource from the [`World`].
@@ -85,45 +84,26 @@ pub(crate) fn internal_apply_state_transition<S: States>(
     current_state: Option<ResMut<State<S>>>,
     new_state: Option<S>,
 ) {
-    match new_state {
-        Some(entered) => {
-            match current_state {
-                // If the [`State<S>`] resource exists, and the state is not the one we are
-                // 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 = mem::replace(&mut state_resource.0, entered.clone());
-
-                        event.send(StateTransitionEvent {
-                            before: Some(exited.clone()),
-                            after: 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()));
-
-                    event.send(StateTransitionEvent {
-                        before: None,
-                        after: Some(entered.clone()),
-                    });
-                }
-            };
-        }
-        None => {
-            // We first remove the [`State<S>`] resource, and if one existed we compute dependant states, send a transition event and run the `OnExit` schedule.
-            if let Some(resource) = current_state {
+    event.send(StateTransitionEvent {
+        before: match (current_state, new_state.as_ref()) {
+            // Update the state resource.
+            (Some(mut current_state), Some(new_state)) => {
+                Some(mem::replace(&mut current_state.0, new_state.clone()))
+            }
+            // Insert the state resource.
+            (None, Some(new_state)) => {
+                commands.insert_resource(State(new_state.clone()));
+                None
+            }
+            // Remove the state resource.
+            (Some(current_state), None) => {
                 commands.remove_resource::<State<S>>();
-
-                event.send(StateTransitionEvent {
-                    before: Some(resource.get().clone()),
-                    after: None,
-                });
+                Some(current_state.get().clone())
             }
-        }
-    }
+            (None, None) => return,
+        },
+        after: new_state,
+    });
 }
 
 /// Sets up the schedules and systems for handling state transitions
@@ -175,31 +155,18 @@ pub fn apply_state_transition<S: FreelyMutableState>(
     next_state: Option<ResMut<NextState<S>>>,
 ) {
     // We want to check if the State and NextState resources exist
-    let Some(mut next_state_resource) = next_state else {
+    let Some(mut next_state) = next_state else {
         return;
     };
+    let next_state = mem::take(next_state.as_mut());
 
-    match next_state_resource.as_ref() {
-        NextState::Pending(new_state) => {
-            if let Some(current_state) = current_state {
-                if new_state != current_state.get() {
-                    let new_state = new_state.clone();
-                    internal_apply_state_transition(
-                        event,
-                        commands,
-                        Some(current_state),
-                        Some(new_state),
-                    );
-                }
-            }
-        }
-        NextState::Unchanged => {
-            // This is the default value, so we don't need to re-insert the resource
-            return;
-        }
+    if current_state.is_none() {
+        return;
     }
 
-    *next_state_resource.as_mut() = NextState::<S>::Unchanged;
+    if let NextState::Pending(next_state) = next_state {
+        internal_apply_state_transition(event, commands, current_state, Some(next_state));
+    }
 }
 
 pub(crate) fn should_run_transition<S: States, T: ScheduleLabel>(

examples/games/alien_cake_addict.rs

@@ -35,6 +35,7 @@ fn main() {
                 rotate_bonus,
                 scoreboard_system,
                 spawn_bonus,
+                restart_keyboard,
             )
                 .run_if(in_state(GameState::Playing)),
         )
@@ -120,6 +121,7 @@ fn setup(mut commands: Commands, asset_server: Res<AssetServer>, mut game: ResMu
     game.player.i = BOARD_SIZE_I / 2;
     game.player.j = BOARD_SIZE_J / 2;
     game.player.move_cooldown = Timer::from_seconds(0.3, TimerMode::Once);
+    game.bonus.entity = None;
 
     commands.spawn(PointLightBundle {
         transform: Transform::from_xyz(4.0, 10.0, 4.0),
@@ -385,7 +387,17 @@ fn scoreboard_system(game: Res<Game>, mut query: Query<&mut Text>) {
     text.sections[0].value = format!("Sugar Rush: {}", game.score);
 }
 
-// restart the game when pressing spacebar
+// restart the game on R press
+fn restart_keyboard(
+    mut next_state: ResMut<NextState<GameState>>,
+    keyboard_input: Res<ButtonInput<KeyCode>>,
+) {
+    if keyboard_input.just_pressed(KeyCode::KeyR) {
+        next_state.set(GameState::Playing);
+    }
+}
+
+// start a new game on spacebar press
 fn gameover_keyboard(
     mut next_state: ResMut<NextState<GameState>>,
     keyboard_input: Res<ButtonInput<KeyCode>>,