Optional state (#11417)

# Objective

Adjust bevy internals to utilize `Option<Res<State<S>>>` instead of
`Res<State<S>>`, to allow for adding/removing states at runtime and
avoid unexpected panics.

As requested here:
#10088 (comment)

---

## Changelog

- Changed the use of `world.resource`/`world.resource_mut` to
`world.get_resource`/`world.get_resource_mut` in the
`run_enter_schedule` and `apply_state_transition` systems and handled
the `None` option.
- `in_state` now returns a ` FnMut(Option<Res<State<S>>>) -> bool +
Clone`, returning `false` if the resource doesn't exist.
- `state_exists_and_equals` was marked as deprecated, and now just runs
and returns `in_state`, since their bevhaviour is now identical
- `state_changed` now takes an `Option<Res<State<S>>>` and returns
`false` if it does not exist.

I would like to remove `state_exists_and_equals` fully, but wanted to
ensure that is acceptable before doing so.

---------

Co-authored-by: Mike <[email protected]>

Author: lee-orr

crates/bevy_ecs/src/schedule/condition.rs

@@ -194,6 +194,8 @@ mod sealed {
 
 /// A collection of [run conditions](Condition) that may be useful in any bevy app.
 pub mod common_conditions {
+    use bevy_utils::warn_once;
+
     use super::NotSystem;
     use crate::{
         change_detection::DetectChanges,
@@ -701,9 +703,7 @@ pub mod common_conditions {
     /// Generates a [`Condition`](super::Condition)-satisfying closure that returns `true`
     /// if the state machine is currently in `state`.
     ///
-    /// # Panics
-    ///
-    /// The condition will panic if the resource does not exist.
+    /// Will return `false` if the state does not exist or if not in `state`.
     ///
     /// # Example
     ///
@@ -748,10 +748,26 @@ pub mod common_conditions {
     /// app.run(&mut world);
     /// assert_eq!(world.resource::<Counter>().0, 0);
     /// ```
-    pub fn in_state<S: States>(state: S) -> impl FnMut(Res<State<S>>) -> bool + Clone {
-        move |current_state: Res<State<S>>| *current_state == state
+    pub fn in_state<S: States>(state: S) -> impl FnMut(Option<Res<State<S>>>) -> bool + Clone {
+        move |current_state: Option<Res<State<S>>>| match current_state {
+            Some(current_state) => *current_state == state,
+            None => {
+                warn_once!("No state matching the type for {} exists - did you forget to `add_state` when initializing the app?", {
+                        let debug_state = format!("{state:?}");
+                        let result = debug_state
+                            .split("::")
+                            .next()
+                            .unwrap_or("Unknown State Type");
+                        result.to_string()
+                    });
+
+                false
+            }
+        }
     }
 
+    /// Identical to [`in_state`] - use that instead.
+    ///
     /// Generates a [`Condition`](super::Condition)-satisfying closure that returns `true`
     /// if the state machine exists and is currently in `state`.
     ///
@@ -804,13 +820,11 @@ pub mod common_conditions {
     /// app.run(&mut world);
     /// assert_eq!(world.resource::<Counter>().0, 0);
     /// ```
+    #[deprecated(since = "0.13.0", note = "use `in_state` instead.")]
     pub fn state_exists_and_equals<S: States>(
         state: S,
     ) -> impl FnMut(Option<Res<State<S>>>) -> bool + Clone {
-        move |current_state: Option<Res<State<S>>>| match current_state {
-            Some(current_state) => *current_state == state,
-            None => false,
-        }
+        in_state(state)
     }
 
     /// A [`Condition`](super::Condition)-satisfying system that returns `true`
@@ -819,9 +833,7 @@ pub mod common_conditions {
     /// To do things on transitions to/from specific states, use their respective OnEnter/OnExit
     /// schedules. Use this run condition if you want to detect any change, regardless of the value.
     ///
-    /// # Panics
-    ///
-    /// The condition will panic if the resource does not exist.
+    /// Returns false if the state does not exist or the state has not changed.
     ///
     /// # Example
     ///
@@ -866,7 +878,10 @@ pub mod common_conditions {
     /// app.run(&mut world);
     /// assert_eq!(world.resource::<Counter>().0, 2);
     /// ```
-    pub fn state_changed<S: States>(current_state: Res<State<S>>) -> bool {
+    pub fn state_changed<S: States>(current_state: Option<Res<State<S>>>) -> bool {
+        let Some(current_state) = current_state else {
+            return false;
+        };
         current_state.is_changed()
     }
 

crates/bevy_ecs/src/schedule/state.rs

@@ -188,9 +188,10 @@ pub struct StateTransitionEvent<S: States> {
 
 /// Run the enter schedule (if it exists) for the current state.
 pub fn run_enter_schedule<S: States>(world: &mut World) {
-    world
-        .try_run_schedule(OnEnter(world.resource::<State<S>>().0.clone()))
-        .ok();
+    let Some(state) = world.get_resource::<State<S>>() else {
+        return;
+    };
+    world.try_run_schedule(OnEnter(state.0.clone())).ok();
 }
 
 /// If a new state is queued in [`NextState<S>`], this system:
@@ -202,26 +203,34 @@ pub fn run_enter_schedule<S: States>(world: &mut World) {
 pub fn apply_state_transition<S: States>(world: &mut World) {
     // We want to take the `NextState` resource,
     // but only mark it as changed if it wasn't empty.
-    let mut next_state_resource = world.resource_mut::<NextState<S>>();
+    let Some(mut next_state_resource) = world.get_resource_mut::<NextState<S>>() else {
+        return;
+    };
     if let Some(entered) = next_state_resource.bypass_change_detection().0.take() {
         next_state_resource.set_changed();
-
-        let mut state_resource = world.resource_mut::<State<S>>();
-        if *state_resource != entered {
-            let exited = mem::replace(&mut state_resource.0, entered.clone());
-            world.send_event(StateTransitionEvent {
-                before: exited.clone(),
-                after: entered.clone(),
-            });
-            // Try to run the schedules if they exist.
-            world.try_run_schedule(OnExit(exited.clone())).ok();
-            world
-                .try_run_schedule(OnTransition {
-                    from: exited,
-                    to: entered.clone(),
-                })
-                .ok();
-            world.try_run_schedule(OnEnter(entered)).ok();
-        }
+        match world.get_resource_mut::<State<S>>() {
+            Some(mut state_resource) => {
+                if *state_resource != entered {
+                    let exited = mem::replace(&mut state_resource.0, entered.clone());
+                    world.send_event(StateTransitionEvent {
+                        before: exited.clone(),
+                        after: entered.clone(),
+                    });
+                    // Try to run the schedules if they exist.
+                    world.try_run_schedule(OnExit(exited.clone())).ok();
+                    world
+                        .try_run_schedule(OnTransition {
+                            from: exited,
+                            to: entered.clone(),
+                        })
+                        .ok();
+                    world.try_run_schedule(OnEnter(entered)).ok();
+                }
+            }
+            None => {
+                world.insert_resource(State(entered.clone()));
+                world.try_run_schedule(OnEnter(entered)).ok();
+            }
+        };
     }
 }