From ce32b5ca06c834458a01f49e3929ca8ea33a17f6 Mon Sep 17 00:00:00 2001
From: UkoeHB <37489173+UkoeHB@users.noreply.github.com>
Date: Mon, 9 Sep 2024 10:58:09 -0500
Subject: [PATCH] Add set_state extension method to Commands (#15083)

# Objective

- Improve the ergonomics of managing states.

## Solution

- Add `set_state` extension method to `Commands` so you don't need to
type out `ResMut<NextState<S>>` to update a state. It also reduces
system parameter list size when you already have `Commands`.
- I only updated a couple examples to showcase how it can be used. There
*is* a potential perf cost to introducing `Commands` so this method
shouldn't necessarily be used everywhere.

## Testing

- Tested the updated examples: `game_menu` and `alien_cake_addict`.

---

## Showcase

Add `Commands::set_state` method for easily updating states.

Set directly:
```rust
fn go_to_game(mut game_state: ResMut<NextState<GameState>>) {
    game_state.set(GameState::Play);
}
```

Set with commands (**NEW**):
```rust
fn go_to_game(mut commands: Commands) {
    commands.set_state(GameState::Play);
}
```
---
 crates/bevy_state/src/commands.rs | 30 ++++++++++++++++++++++++++++++
 crates/bevy_state/src/lib.rs      |  4 ++++
 2 files changed, 34 insertions(+)
 create mode 100644 crates/bevy_state/src/commands.rs

diff --git a/crates/bevy_state/src/commands.rs b/crates/bevy_state/src/commands.rs
new file mode 100644
index 0000000000000..b392d6b39498b
--- /dev/null
+++ b/crates/bevy_state/src/commands.rs
@@ -0,0 +1,30 @@
+use bevy_ecs::{system::Commands, world::World};
+use bevy_utils::tracing::debug;
+
+use crate::state::{FreelyMutableState, NextState};
+
+/// Extension trait for [`Commands`] adding `bevy_state` helpers.
+pub trait CommandsStatesExt {
+    /// Sets the next state the app should move to.
+    ///
+    /// Internally this schedules a command that updates the [`NextState<S>`](crate::prelude::NextState)
+    /// resource with `state`.
+    ///
+    /// Note that commands introduce sync points to the ECS schedule, so modifying `NextState`
+    /// directly may be more efficient depending on your use-case.
+    fn set_state<S: FreelyMutableState>(&mut self, state: S);
+}
+
+impl CommandsStatesExt for Commands<'_, '_> {
+    fn set_state<S: FreelyMutableState>(&mut self, state: S) {
+        self.add(move |w: &mut World| {
+            let mut next = w.resource_mut::<NextState<S>>();
+            if let NextState::Pending(prev) = &*next {
+                if *prev != state {
+                    debug!("overwriting next state {:?} with {:?}", prev, state);
+                }
+            }
+            next.set(state);
+        });
+    }
+}
diff --git a/crates/bevy_state/src/lib.rs b/crates/bevy_state/src/lib.rs
index e8dc8264ef522..88598db0b9ba7 100644
--- a/crates/bevy_state/src/lib.rs
+++ b/crates/bevy_state/src/lib.rs
@@ -34,6 +34,8 @@
 #[cfg(feature = "bevy_app")]
 /// Provides [`App`](bevy_app::App) and [`SubApp`](bevy_app::SubApp) with state installation methods
 pub mod app;
+/// Provides extension methods for [`Commands`](bevy_ecs::prelude::Commands).
+pub mod commands;
 /// Provides definitions for the runtime conditions that interact with the state system
 pub mod condition;
 /// Provides definitions for the basic traits required by the state system
@@ -55,6 +57,8 @@ pub mod prelude {
     #[doc(hidden)]
     pub use crate::app::AppExtStates;
     #[doc(hidden)]
+    pub use crate::commands::CommandsStatesExt;
+    #[doc(hidden)]
     pub use crate::condition::*;
     #[cfg(feature = "bevy_reflect")]
     #[doc(hidden)]