Computed State & Sub States (#11426)

## Summary/Description
This PR extends states to allow support for a wider variety of state
types and patterns, by providing 3 distinct types of state:
- Standard [`States`] can only be changed by manually setting the
[`NextState<S>`] resource. These states are the baseline on which the
other state types are built, and can be used on their own for many
simple patterns. See the [state
example](https://github.com/bevyengine/bevy/blob/latest/examples/ecs/state.rs)
for a simple use case - these are the states that existed so far in
Bevy.
- [`SubStates`] are children of other states - they can be changed
manually using [`NextState<S>`], but are removed from the [`World`] if
the source states aren't in the right state. See the [sub_states
example](https://github.com/lee-orr/bevy/blob/derived_state/examples/ecs/sub_states.rs)
for a simple use case based on the derive macro, or read the trait docs
for more complex scenarios.
- [`ComputedStates`] are fully derived from other states - they provide
a [`compute`](ComputedStates::compute) method that takes in the source
states and returns their derived value. They are particularly useful for
situations where a simplified view of the source states is necessary -
such as having an `InAMenu` computed state derived from a source state
that defines multiple distinct menus. See the [computed state
example](https://github.com/lee-orr/bevy/blob/derived_state/examples/ecs/computed_states.rscomputed_states.rs)
to see a sampling of uses for these states.

# Objective

This PR is another attempt at allowing Bevy to better handle complex
state objects in a manner that doesn't rely on strict equality. While my
previous attempts (#10088 and
#9957) relied on complex matching
capacities at the point of adding a system to application, this one
instead relies on deterministically deriving simple states from more
complex ones.

As a result, it does not require any special macros, nor does it change
any other interactions with the state system once you define and add
your derived state. It also maintains a degree of distinction between
`State` and just normal application state - your derivations have to end
up being discreet pre-determined values, meaning there is less of a
risk/temptation to place a significant amount of logic and data within a
given state.

### Addition - Sub States
closes #9942 
After some conversation with Maintainers & SMEs, a significant concern
was that people might attempt to use this feature as if it were
sub-states, and find themselves unable to use it appropriately. Since
`ComputedState` is mainly a state matching feature, while `SubStates`
are more of a state mutation related feature - but one that is easy to
add with the help of the machinery introduced by `ComputedState`, it was
added here as well. The relevant discussion is here:
https://discord.com/channels/691052431525675048/1200556329803186316

## Solution
closes #11358 

The solution is to create a new type of state - one implementing
`ComputedStates` - which is deterministically tied to one or more other
states. Implementors write a function to transform the source states
into the computed state, and it gets triggered whenever one of the
source states changes.

In addition, we added the `FreelyMutableState` trait , which is
implemented as part of the derive macro for `States`. This allows us to
limit use of `NextState<S>` to states that are actually mutable,
preventing mis-use of `ComputedStates`.

---

## Changelog

- Added `ComputedStates` trait
- Added `FreelyMutableState` trait
- Converted `NextState` resource to an Enum, with `Unchanged` and
`Pending`
- Added `App::add_computed_state::<S: ComputedStates>()`, to allow for
easily adding derived states to an App.
- Moved the `StateTransition` schedule label from `bevy_app` to
`bevy_ecs` - but maintained the export in `bevy_app` for continuity.
- Modified the process for updating states. Instead of just having an
`apply_state_transition` system that can be added anywhere, we now have
a multi-stage process that has to run within the `StateTransition`
label. First, all the state changes are calculated - manual transitions
rely on `apply_state_transition`, while computed transitions run their
computation process before both call `internal_apply_state_transition`
to apply the transition, send out the transition event, trigger
dependent states, and record which exit/transition/enter schedules need
to occur. Once all the states have been updated, the transition
schedules are called - first the exit schedules, then transition
schedules and finally enter schedules.
- Added `SubStates` trait
- Adjusted `apply_state_transition` to be a no-op if the `State<S>`
resource doesn't exist

## Migration Guide

If the user accessed the NextState resource's value directly or created
them from scratch they will need to adjust to use the new enum variants:
- if they created a `NextState(Some(S))` - they should now use
`NextState::Pending(S)`
- if they created a `NextState(None)` -they should now use
`NextState::Unchanged`
- if they matched on the `NextState` value, they would need to make the
adjustments above

If the user manually utilized `apply_state_transition`, they should
instead use systems that trigger the `StateTransition` schedule.

---
## Future Work
There is still some future potential work in the area, but I wanted to
keep these potential features and changes separate to keep the scope
here contained, and keep the core of it easy to understand and use.
However, I do want to note some of these things, both as inspiration to
others and an illustration of what this PR could unlock.

- `NextState::Remove` - Now that the `State` related mechanisms all
utilize options (#11417), it's fairly easy to add support for explicit
state removal. And while `ComputedStates` can add and remove themselves,
right now `FreelyMutableState`s can't be removed from within the state
system. While it existed originally in this PR, it is a different
question with a separate scope and usability concerns - so having it as
it's own future PR seems like the best approach. This feature currently
lives in a separate branch in my fork, and the differences between it
and this PR can be seen here: lee-orr#5

- `NextState::ReEnter` - this would allow you to trigger exit & entry
systems for the current state type. We can potentially also add a
`NextState::ReEnterRecirsive` to also re-trigger any states that depend
on the current one.

- More mechanisms for `State` updates - This PR would finally make
states that aren't a set of exclusive Enums useful, and with that comes
the question of setting state more effectively. Right now, to update a
state you either need to fully create the new state, or include the
`Res<Option<State<S>>>` resource in your system, clone the state, mutate
it, and then use `NextState.set(my_mutated_state)` to make it the
pending next state. There are a few other potential methods that could
be implemented in future PRs:
- Inverse Compute States - these would essentially be compute states
that have an additional (manually defined) function that can be used to
nudge the source states so that they result in the computed states
having a given value. For example, you could use set the `IsPaused`
state, and it would attempt to pause or unpause the game by modifying
the `AppState` as needed.
- Closure-based state modification - this would involve adding a
`NextState.modify(f: impl Fn(Option<S> -> Option<S>)` method, and then
you can pass in closures or function pointers to adjust the state as
needed.
- Message-based state modification - this would involve either creating
states that can respond to specific messages, similar to Elm or Redux.
These could either use the `NextState` mechanism or the Event mechanism.

- ~`SubStates` - which are essentially a hybrid of computed and manual
states. In the simplest (and most likely) version, they would work by
having a computed element that determines whether the state should
exist, and if it should has the capacity to add a new version in, but
then any changes to it's content would be freely mutated.~ this feature
is now part of this PR. See above.

- Lastly, since states are getting more complex there might be value in
moving them out of `bevy_ecs` and into their own crate, or at least out
of the `schedule` module into a `states` module. #11087

As mentioned, all these future work elements are TBD and are explicitly
not part of this PR - I just wanted to provide them as potential
explorations for the future.

---------

Co-authored-by: Alice Cecile <alice.i.cecile@gmail.com>
Co-authored-by: Marcel Champagne <voiceofmarcel@gmail.com>
Co-authored-by: MiniaczQ <xnetroidpl@gmail.com>

Cargo.toml

@@ -1692,6 +1692,28 @@ description = "Illustrates how to use States to control transitioning from a Men
 category = "ECS (Entity Component System)"
 wasm = false
 
+[[example]]
+name = "sub_states"
+path = "examples/ecs/sub_states.rs"
+doc-scrape-examples = true
+
+[package.metadata.example.sub_states]
+name = "Sub States"
+description = "Using Sub States for hierarchical state handling."
+category = "ECS (Entity Component System)"
+wasm = false
+
+[[example]]
+name = "computed_states"
+path = "examples/ecs/computed_states.rs"
+doc-scrape-examples = true
+
+[package.metadata.example.computed_states]
+name = "Computed States"
+description = "Advanced state patterns using Computed States"
+category = "ECS (Entity Component System)"
+wasm = false
+
 [[example]]
 name = "system_piping"
 path = "examples/ecs/system_piping.rs"

crates/bevy_app/src/app.rs

@@ -7,7 +7,7 @@ use bevy_ecs::{
     event::{event_update_system, ManualEventReader},
     intern::Interned,
     prelude::*,
-    schedule::{ScheduleBuildSettings, ScheduleLabel},
+    schedule::{FreelyMutableState, ScheduleBuildSettings, ScheduleLabel},
     system::SystemId,
 };
 #[cfg(feature = "trace")]
@@ -266,56 +266,60 @@ impl App {
 
     /// Initializes a [`State`] with standard starting values.
     ///
-    /// If the [`State`] already exists, nothing happens.
+    /// This method is idempotent: it has no effect when called again using the same generic type.
     ///
-    /// Adds [`State<S>`] and [`NextState<S>`] resources, [`OnEnter`] and [`OnExit`] schedules for
-    /// each state variant (if they don't already exist), an instance of [`apply_state_transition::<S>`]
-    /// in [`StateTransition`] so that transitions happen before [`Update`] and an instance of
-    /// [`run_enter_schedule::<S>`] in [`StateTransition`] with a [`run_once`] condition to run the
-    /// on enter schedule of the initial state.
+    /// Adds [`State<S>`] and [`NextState<S>`] resources, and enables use of the [`OnEnter`], [`OnTransition`] and [`OnExit`] schedules.
+    /// These schedules are triggered before [`Update`](crate::Update) and at startup.
     ///
     /// If you would like to control how other systems run based on the current state, you can
     /// emulate this behavior using the [`in_state`] [`Condition`].
     ///
-    /// Note that you can also apply state transitions at other points in the schedule by adding
-    /// the [`apply_state_transition::<S>`] system manually.
-    ///
-    /// [`StateTransition`]: crate::StateTransition
-    /// [`Update`]: crate::Update
-    /// [`run_once`]: bevy_ecs::schedule::common_conditions::run_once
-    /// [`run_enter_schedule::<S>`]: bevy_ecs::schedule::run_enter_schedule
-    /// [`apply_state_transition::<S>`]: bevy_ecs::schedule::apply_state_transition
-    pub fn init_state<S: States + FromWorld>(&mut self) -> &mut Self {
+    /// Note that you can also apply state transitions at other points in the schedule
+    /// by triggering the [`StateTransition`](`bevy_ecs::schedule::StateTransition`) schedule manually.
+    pub fn init_state<S: FreelyMutableState + FromWorld>(&mut self) -> &mut Self {
         self.main_mut().init_state::<S>();
         self
     }
 
     /// Inserts a specific [`State`] to the current [`App`] and overrides any [`State`] previously
     /// added of the same type.
     ///
-    /// Adds [`State<S>`] and [`NextState<S>`] resources, [`OnEnter`] and [`OnExit`] schedules for
-    /// each state variant (if they don't already exist), an instance of [`apply_state_transition::<S>`]
-    /// in [`StateTransition`] so that transitions happen before [`Update`](crate::Update) and an
-    /// instance of [`run_enter_schedule::<S>`] in [`StateTransition`] with a [`run_once`]
-    /// condition to run the on enter schedule of the initial state.
+    /// Adds [`State<S>`] and [`NextState<S>`] resources, and enables use of the [`OnEnter`], [`OnTransition`] and [`OnExit`] schedules.
+    /// These schedules are triggered before [`Update`](crate::Update) and at startup.
     ///
     /// If you would like to control how other systems run based on the current state, you can
     /// emulate this behavior using the [`in_state`] [`Condition`].
     ///
-    /// Note that you can also apply state transitions at other points in the schedule by adding
-    /// the [`apply_state_transition::<S>`] system manually.
+    /// Note that you can also apply state transitions at other points in the schedule
+    /// by triggering the [`StateTransition`](`bevy_ecs::schedule::StateTransition`) schedule manually.
+    pub fn insert_state<S: FreelyMutableState>(&mut self, state: S) -> &mut Self {
+        self.main_mut().insert_state::<S>(state);
+        self
+    }
+
+    /// Sets up a type implementing [`ComputedStates`].
+    ///
+    /// This method is idempotent: it has no effect when called again using the same generic type.
+    ///
+    /// For each source state the derived state depends on, it adds this state's derivation
+    /// to it's [`ComputeDependantStates<Source>`](bevy_ecs::schedule::ComputeDependantStates<S>) schedule.
+    pub fn add_computed_state<S: ComputedStates>(&mut self) -> &mut Self {
+        self.main_mut().add_computed_state::<S>();
+        self
+    }
+
+    /// Sets up a type implementing [`SubStates`].
+    ///
+    /// This method is idempotent: it has no effect when called again using the same generic type.
     ///
-    /// [`StateTransition`]: crate::StateTransition
-    /// [`Update`]: crate::Update
-    /// [`run_once`]: bevy_ecs::schedule::common_conditions::run_once
-    /// [`run_enter_schedule::<S>`]: bevy_ecs::schedule::run_enter_schedule
-    /// [`apply_state_transition::<S>`]: bevy_ecs::schedule::apply_state_transition
-    pub fn insert_state<S: States>(&mut self, state: S) -> &mut Self {
-        self.main_mut().insert_state(state);
+    /// For each source state the derived state depends on, it adds this state's existence check
+    /// to it's [`ComputeDependantStates<Source>`](bevy_ecs::schedule::ComputeDependantStates<S>) schedule.
+    pub fn add_sub_state<S: SubStates>(&mut self) -> &mut Self {
+        self.main_mut().add_sub_state::<S>();
         self
     }
 
-    /// Adds a collection of systems to `schedule` (stored in the main world's [`Schedules`]).
+    /// Adds one or more systems to the given schedule in this app's [`Schedules`].
     ///
     /// # Examples
     ///

crates/bevy_app/src/lib.rs

@@ -31,8 +31,7 @@ pub mod prelude {
         app::{App, AppExit},
         main_schedule::{
             First, FixedFirst, FixedLast, FixedPostUpdate, FixedPreUpdate, FixedUpdate, Last, Main,
-            PostStartup, PostUpdate, PreStartup, PreUpdate, SpawnScene, Startup, StateTransition,
-            Update,
+            PostStartup, PostUpdate, PreStartup, PreUpdate, SpawnScene, Startup, Update,
         },
         sub_app::SubApp,
         DynamicPlugin, Plugin, PluginGroup,

crates/bevy_app/src/main_schedule.rs

@@ -1,6 +1,6 @@
 use crate::{App, Plugin};
 use bevy_ecs::{
-    schedule::{ExecutorKind, InternedScheduleLabel, Schedule, ScheduleLabel},
+    schedule::{ExecutorKind, InternedScheduleLabel, Schedule, ScheduleLabel, StateTransition},
     system::{Local, Resource},
     world::{Mut, World},
 };
@@ -73,12 +73,6 @@ pub struct First;
 #[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
 pub struct PreUpdate;
 
-/// Runs [state transitions](bevy_ecs::schedule::States).
-///
-/// See the [`Main`] schedule for some details about how schedules are run.
-#[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
-pub struct StateTransition;
-
 /// Runs the [`FixedMain`] schedule in a loop according until all relevant elapsed time has been "consumed".
 ///
 /// See the [`Main`] schedule for some details about how schedules are run.

crates/bevy_app/src/sub_app.rs

@@ -1,10 +1,10 @@
-use crate::{App, InternedAppLabel, Plugin, Plugins, PluginsState, StateTransition};
+use crate::{App, InternedAppLabel, Plugin, Plugins, PluginsState, Startup};
 use bevy_ecs::{
     event::EventRegistry,
     prelude::*,
     schedule::{
-        common_conditions::run_once as run_once_condition, run_enter_schedule,
-        InternedScheduleLabel, ScheduleBuildSettings, ScheduleLabel,
+        setup_state_transitions_in_world, FreelyMutableState, InternedScheduleLabel,
+        ScheduleBuildSettings, ScheduleLabel,
     },
     system::SystemId,
 };
@@ -317,40 +317,61 @@ impl SubApp {
     }
 
     /// See [`App::init_state`].
-    pub fn init_state<S: States + FromWorld>(&mut self) -> &mut Self {
+    pub fn init_state<S: FreelyMutableState + FromWorld>(&mut self) -> &mut Self {
         if !self.world.contains_resource::<State<S>>() {
+            setup_state_transitions_in_world(&mut self.world, Some(Startup.intern()));
             self.init_resource::<State<S>>()
                 .init_resource::<NextState<S>>()
-                .add_event::<StateTransitionEvent<S>>()
-                .add_systems(
-                    StateTransition,
-                    (
-                        run_enter_schedule::<S>.run_if(run_once_condition()),
-                        apply_state_transition::<S>,
-                    )
-                        .chain(),
-                );
+                .add_event::<StateTransitionEvent<S>>();
+            let schedule = self.get_schedule_mut(StateTransition).unwrap();
+            S::register_state(schedule);
         }
 
-        // The OnEnter, OnExit, and OnTransition schedules are lazily initialized
-        // (i.e. when the first system is added to them), so World::try_run_schedule
-        // is used to fail gracefully if they aren't present.
         self
     }
 
     /// See [`App::insert_state`].
-    pub fn insert_state<S: States>(&mut self, state: S) -> &mut Self {
-        self.insert_resource(State::new(state))
-            .init_resource::<NextState<S>>()
-            .add_event::<StateTransitionEvent<S>>()
-            .add_systems(
-                StateTransition,
-                (
-                    run_enter_schedule::<S>.run_if(run_once_condition()),
-                    apply_state_transition::<S>,
-                )
-                    .chain(),
-            );
+    pub fn insert_state<S: FreelyMutableState>(&mut self, state: S) -> &mut Self {
+        if !self.world.contains_resource::<State<S>>() {
+            setup_state_transitions_in_world(&mut self.world, Some(Startup.intern()));
+            self.insert_resource::<State<S>>(State::new(state))
+                .init_resource::<NextState<S>>()
+                .add_event::<StateTransitionEvent<S>>();
+
+            let schedule = self.get_schedule_mut(StateTransition).unwrap();
+            S::register_state(schedule);
+        }
+
+        self
+    }
+
+    /// See [`App::add_computed_state`].
+    pub fn add_computed_state<S: ComputedStates>(&mut self) -> &mut Self {
+        if !self
+            .world
+            .contains_resource::<Events<StateTransitionEvent<S>>>()
+        {
+            setup_state_transitions_in_world(&mut self.world, Some(Startup.intern()));
+            self.add_event::<StateTransitionEvent<S>>();
+            let schedule = self.get_schedule_mut(StateTransition).unwrap();
+            S::register_computed_state_systems(schedule);
+        }
+
+        self
+    }
+
+    /// See [`App::add_sub_state`].
+    pub fn add_sub_state<S: SubStates>(&mut self) -> &mut Self {
+        if !self
+            .world
+            .contains_resource::<Events<StateTransitionEvent<S>>>()
+        {
+            setup_state_transitions_in_world(&mut self.world, Some(Startup.intern()));
+            self.init_resource::<NextState<S>>();
+            self.add_event::<StateTransitionEvent<S>>();
+            let schedule = self.get_schedule_mut(StateTransition).unwrap();
+            S::register_sub_state_systems(schedule);
+        }
 
         self
     }

crates/bevy_ecs/macros/Cargo.toml

@@ -11,7 +11,7 @@ proc-macro = true
 [dependencies]
 bevy_macro_utils = { path = "../../bevy_macro_utils", version = "0.14.0-dev" }
 
-syn = "2.0"
+syn = { version = "2.0", features = ["full"] }
 quote = "1.0"
 proc-macro2 = "1.0"
 

crates/bevy_ecs/macros/src/lib.rs

@@ -521,3 +521,8 @@ pub fn derive_component(input: TokenStream) -> TokenStream {
 pub fn derive_states(input: TokenStream) -> TokenStream {
     states::derive_states(input)
 }
+
+#[proc_macro_derive(SubStates, attributes(source))]
+pub fn derive_substates(input: TokenStream) -> TokenStream {
+    states::derive_substates(input)
+}