Add more granular system sets for state transition schedule ordering (#13763)

# Objective

Fixes #13711 

## Solution

Introduce smaller, generic system sets for each schedule variant, which
are ordered against other generic variants:
- `ExitSchedules<S>` - For `OnExit` schedules, runs from leaf states to
root states.
- `TransitionSchedules<S>` - For `OnTransition` schedules, runs in
arbitrary order.
- `EnterSchedules<S>` - For `OnEnter` schedules, runs from root states
to leaf states.

Also unified `ApplyStateTransition<S>` schedule which works in basically
the same way, just for internals.

## Testing

- One test that tests schedule execution order

---------

Co-authored-by: Lee-Orr <lee-orr@users.noreply.github.com>

crates/bevy_state/src/lib.rs

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

crates/bevy_state/src/state/freely_mutable_state.rs

@@ -17,27 +17,33 @@ use super::{take_next_state, transitions::*};
 pub trait FreelyMutableState: States {
     /// This function registers all the necessary systems to apply state changes and run transition schedules
     fn register_state(schedule: &mut Schedule) {
+        schedule.configure_sets((
+            ApplyStateTransition::<Self>::default()
+                .in_set(StateTransitionSteps::DependentTransitions),
+            ExitSchedules::<Self>::default().in_set(StateTransitionSteps::ExitSchedules),
+            TransitionSchedules::<Self>::default()
+                .in_set(StateTransitionSteps::TransitionSchedules),
+            EnterSchedules::<Self>::default().in_set(StateTransitionSteps::EnterSchedules),
+        ));
+
         schedule
             .add_systems(
-                apply_state_transition::<Self>.in_set(ApplyStateTransition::<Self>::apply()),
+                apply_state_transition::<Self>.in_set(ApplyStateTransition::<Self>::default()),
-            )
-            .add_systems(
-                last_transition::<Self>
-                    .pipe(run_enter::<Self>)
-                    .in_set(StateTransitionSteps::EnterSchedules),
             )
             .add_systems(
                 last_transition::<Self>
                     .pipe(run_exit::<Self>)
-                    .in_set(StateTransitionSteps::ExitSchedules),
+                    .in_set(ExitSchedules::<Self>::default()),
             )
             .add_systems(
                 last_transition::<Self>
                     .pipe(run_transition::<Self>)
-                    .in_set(StateTransitionSteps::TransitionSchedules),
+                    .in_set(TransitionSchedules::<Self>::default()),
             )
-            .configure_sets(
+            .add_systems(
-                ApplyStateTransition::<Self>::apply().in_set(StateTransitionSteps::RootTransitions),
+                last_transition::<Self>
+                    .pipe(run_enter::<Self>)
+                    .in_set(EnterSchedules::<Self>::default()),
             );
     }
 }

crates/bevy_state/src/state/mod.rs

@@ -508,14 +508,13 @@ mod tests {
     #[test]
     fn same_state_transition_should_emit_event_and_not_run_schedules() {
         let mut world = World::new();
+        setup_state_transitions_in_world(&mut world, None);
         EventRegistry::register_event::<StateTransitionEvent<SimpleState>>(&mut world);
         world.init_resource::<State<SimpleState>>();
-        let mut schedules = Schedules::new();
+        let mut schedules = world.resource_mut::<Schedules>();
-        let mut apply_changes = Schedule::new(StateTransition);
+        let apply_changes = schedules.get_mut(StateTransition).unwrap();
-        SimpleState::register_state(&mut apply_changes);
+        SimpleState::register_state(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);
@@ -528,9 +527,7 @@ mod tests {
         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(TransitionCounter::default());
-        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);
@@ -546,7 +543,7 @@ mod tests {
             *world.resource::<TransitionCounter>(),
             TransitionCounter {
                 exit: 0,
-                transition: 0,
+                transition: 1, // Same state transitions are allowed
                 enter: 0
             }
         );
@@ -619,4 +616,118 @@ mod tests {
             1
         );
     }
+
+    #[derive(Resource, Default, Debug)]
+    struct TransitionTracker(Vec<&'static str>);
+
+    #[derive(PartialEq, Eq, Debug, Hash, Clone)]
+    enum TransitionTestingComputedState {
+        IsA,
+        IsBAndEven,
+        IsBAndOdd,
+    }
+
+    impl ComputedStates for TransitionTestingComputedState {
+        type SourceStates = (Option<SimpleState>, Option<SubState>);
+
+        fn compute(sources: (Option<SimpleState>, Option<SubState>)) -> Option<Self> {
+            match sources {
+                (Some(simple), sub) => {
+                    if simple == SimpleState::A {
+                        Some(Self::IsA)
+                    } else if sub == Some(SubState::One) {
+                        Some(Self::IsBAndOdd)
+                    } else if sub == Some(SubState::Two) {
+                        Some(Self::IsBAndEven)
+                    } else {
+                        None
+                    }
+                }
+                _ => None,
+            }
+        }
+    }
+
+    #[test]
+    fn check_transition_orders() {
+        let mut world = World::new();
+        setup_state_transitions_in_world(&mut world, None);
+        EventRegistry::register_event::<StateTransitionEvent<SimpleState>>(&mut world);
+        EventRegistry::register_event::<StateTransitionEvent<SubState>>(&mut world);
+        EventRegistry::register_event::<StateTransitionEvent<TransitionTestingComputedState>>(
+            &mut world,
+        );
+        world.insert_resource(State(SimpleState::B(true)));
+        world.init_resource::<State<SubState>>();
+        world.insert_resource(State(TransitionTestingComputedState::IsA));
+        let mut schedules = world.remove_resource::<Schedules>().unwrap();
+        let apply_changes = schedules.get_mut(StateTransition).unwrap();
+        SimpleState::register_state(apply_changes);
+        SubState::register_sub_state_systems(apply_changes);
+        TransitionTestingComputedState::register_computed_state_systems(apply_changes);
+
+        world.init_resource::<TransitionTracker>();
+        fn register_transition(string: &'static str) -> impl Fn(ResMut<TransitionTracker>) {
+            move |mut transitions: ResMut<TransitionTracker>| transitions.0.push(string)
+        }
+
+        schedules.add_systems(
+            StateTransition,
+            register_transition("simple exit").in_set(ExitSchedules::<SimpleState>::default()),
+        );
+        schedules.add_systems(
+            StateTransition,
+            register_transition("simple transition")
+                .in_set(TransitionSchedules::<SimpleState>::default()),
+        );
+        schedules.add_systems(
+            StateTransition,
+            register_transition("simple enter").in_set(EnterSchedules::<SimpleState>::default()),
+        );
+
+        schedules.add_systems(
+            StateTransition,
+            register_transition("sub exit").in_set(ExitSchedules::<SubState>::default()),
+        );
+        schedules.add_systems(
+            StateTransition,
+            register_transition("sub transition")
+                .in_set(TransitionSchedules::<SubState>::default()),
+        );
+        schedules.add_systems(
+            StateTransition,
+            register_transition("sub enter").in_set(EnterSchedules::<SubState>::default()),
+        );
+
+        schedules.add_systems(
+            StateTransition,
+            register_transition("computed exit")
+                .in_set(ExitSchedules::<TransitionTestingComputedState>::default()),
+        );
+        schedules.add_systems(
+            StateTransition,
+            register_transition("computed transition")
+                .in_set(TransitionSchedules::<TransitionTestingComputedState>::default()),
+        );
+        schedules.add_systems(
+            StateTransition,
+            register_transition("computed enter")
+                .in_set(EnterSchedules::<TransitionTestingComputedState>::default()),
+        );
+
+        world.insert_resource(schedules);
+
+        world.run_schedule(StateTransition);
+
+        let transitions = &world.resource::<TransitionTracker>().0;
+
+        assert_eq!(transitions.len(), 9);
+        assert_eq!(transitions[0], "computed exit");
+        assert_eq!(transitions[1], "sub exit");
+        assert_eq!(transitions[2], "simple exit");
+        // Transition order is arbitrary and doesn't need testing.
+        assert_eq!(transitions[6], "simple enter");
+        assert_eq!(transitions[7], "sub enter");
+        assert_eq!(transitions[8], "computed enter");
+    }
 }

crates/bevy_state/src/state/state_set.rs

@@ -10,7 +10,8 @@ use self::sealed::StateSetSealed;
 use super::{
     computed_states::ComputedStates, internal_apply_state_transition, last_transition, run_enter,
     run_exit, run_transition, sub_states::SubStates, take_next_state, ApplyStateTransition,
-    NextState, State, StateTransitionEvent, StateTransitionSteps, States,
+    EnterSchedules, ExitSchedules, NextState, State, StateTransitionEvent, StateTransitionSteps,
+    States, TransitionSchedules,
 };
 
 mod sealed {
@@ -114,27 +115,35 @@ impl<S: InnerStateSet> StateSet for S {
                 internal_apply_state_transition(event, commands, current_state, new_state);
             };
 
+        schedule.configure_sets((
+            ApplyStateTransition::<T>::default()
+                .in_set(StateTransitionSteps::DependentTransitions)
+                .after(ApplyStateTransition::<S::RawState>::default()),
+            ExitSchedules::<T>::default()
+                .in_set(StateTransitionSteps::ExitSchedules)
+                .before(ExitSchedules::<S::RawState>::default()),
+            TransitionSchedules::<T>::default().in_set(StateTransitionSteps::TransitionSchedules),
+            EnterSchedules::<T>::default()
+                .in_set(StateTransitionSteps::EnterSchedules)
+                .after(EnterSchedules::<S::RawState>::default()),
+        ));
+
         schedule
-            .add_systems(apply_state_transition.in_set(ApplyStateTransition::<T>::apply()))
+            .add_systems(apply_state_transition.in_set(ApplyStateTransition::<T>::default()))
-            .add_systems(
-                last_transition::<T>
-                    .pipe(run_enter::<T>)
-                    .in_set(StateTransitionSteps::EnterSchedules),
-            )
             .add_systems(
                 last_transition::<T>
                     .pipe(run_exit::<T>)
-                    .in_set(StateTransitionSteps::ExitSchedules),
+                    .in_set(ExitSchedules::<T>::default()),
             )
             .add_systems(
                 last_transition::<T>
                     .pipe(run_transition::<T>)
-                    .in_set(StateTransitionSteps::TransitionSchedules),
+                    .in_set(TransitionSchedules::<T>::default()),
             )
-            .configure_sets(
+            .add_systems(
-                ApplyStateTransition::<T>::apply()
+                last_transition::<T>
-                    .in_set(StateTransitionSteps::DependentTransitions)
+                    .pipe(run_enter::<T>)
-                    .after(ApplyStateTransition::<S::RawState>::apply()),
+                    .in_set(EnterSchedules::<T>::default()),
             );
     }
 
@@ -186,27 +195,35 @@ impl<S: InnerStateSet> StateSet for S {
                 internal_apply_state_transition(event, commands, current_state_res, new_state);
             };
 
+        schedule.configure_sets((
+            ApplyStateTransition::<T>::default()
+                .in_set(StateTransitionSteps::DependentTransitions)
+                .after(ApplyStateTransition::<S::RawState>::default()),
+            ExitSchedules::<T>::default()
+                .in_set(StateTransitionSteps::ExitSchedules)
+                .before(ExitSchedules::<S::RawState>::default()),
+            TransitionSchedules::<T>::default().in_set(StateTransitionSteps::TransitionSchedules),
+            EnterSchedules::<T>::default()
+                .in_set(StateTransitionSteps::EnterSchedules)
+                .after(EnterSchedules::<S::RawState>::default()),
+        ));
+
         schedule
-            .add_systems(apply_state_transition.in_set(ApplyStateTransition::<T>::apply()))
+            .add_systems(apply_state_transition.in_set(ApplyStateTransition::<T>::default()))
-            .add_systems(
-                last_transition::<T>
-                    .pipe(run_enter::<T>)
-                    .in_set(StateTransitionSteps::EnterSchedules),
-            )
             .add_systems(
                 last_transition::<T>
                     .pipe(run_exit::<T>)
-                    .in_set(StateTransitionSteps::ExitSchedules),
+                    .in_set(ExitSchedules::<T>::default()),
             )
             .add_systems(
                 last_transition::<T>
                     .pipe(run_transition::<T>)
-                    .in_set(StateTransitionSteps::TransitionSchedules),
+                    .in_set(TransitionSchedules::<T>::default()),
             )
-            .configure_sets(
+            .add_systems(
-                ApplyStateTransition::<T>::apply()
+                last_transition::<T>
-                    .in_set(StateTransitionSteps::DependentTransitions)
+                    .pipe(run_enter::<T>)
-                    .after(ApplyStateTransition::<S::RawState>::apply()),
+                    .in_set(EnterSchedules::<T>::default()),
             );
     }
 }
@@ -243,16 +260,25 @@ macro_rules! impl_state_set_sealed_tuples {
                         internal_apply_state_transition(event, commands, current_state, new_state);
                     };
 
-                schedule
+                schedule.configure_sets((
-                    .add_systems(apply_state_transition.in_set(ApplyStateTransition::<T>::apply()))
+                    ApplyStateTransition::<T>::default()
-                    .add_systems(last_transition::<T>.pipe(run_enter::<T>).in_set(StateTransitionSteps::EnterSchedules))
-                    .add_systems(last_transition::<T>.pipe(run_exit::<T>).in_set(StateTransitionSteps::ExitSchedules))
-                    .add_systems(last_transition::<T>.pipe(run_transition::<T>).in_set(StateTransitionSteps::TransitionSchedules))
-                    .configure_sets(
-                        ApplyStateTransition::<T>::apply()
                         .in_set(StateTransitionSteps::DependentTransitions)
-                        $(.after(ApplyStateTransition::<$param::RawState>::apply()))*
+                        $(.after(ApplyStateTransition::<$param::RawState>::default()))*,
-                    );
+                    ExitSchedules::<T>::default()
+                        .in_set(StateTransitionSteps::ExitSchedules)
+                        $(.before(ExitSchedules::<$param::RawState>::default()))*,
+                    TransitionSchedules::<T>::default()
+                        .in_set(StateTransitionSteps::TransitionSchedules),
+                    EnterSchedules::<T>::default()
+                        .in_set(StateTransitionSteps::EnterSchedules)
+                        $(.after(EnterSchedules::<$param::RawState>::default()))*,
+                ));
+
+                schedule
+                    .add_systems(apply_state_transition.in_set(ApplyStateTransition::<T>::default()))
+                    .add_systems(last_transition::<T>.pipe(run_exit::<T>).in_set(ExitSchedules::<T>::default()))
+                    .add_systems(last_transition::<T>.pipe(run_transition::<T>).in_set(TransitionSchedules::<T>::default()))
+                    .add_systems(last_transition::<T>.pipe(run_enter::<T>).in_set(EnterSchedules::<T>::default()));
             }
 
             fn register_sub_state_systems_in_schedule<T: SubStates<SourceStates = Self>>(
@@ -288,16 +314,25 @@ macro_rules! impl_state_set_sealed_tuples {
                         internal_apply_state_transition(event, commands, current_state_res, new_state);
                     };
 
-                schedule
+                schedule.configure_sets((
-                    .add_systems(apply_state_transition.in_set(ApplyStateTransition::<T>::apply()))
+                    ApplyStateTransition::<T>::default()
-                    .add_systems(last_transition::<T>.pipe(run_enter::<T>).in_set(StateTransitionSteps::EnterSchedules))
-                    .add_systems(last_transition::<T>.pipe(run_exit::<T>).in_set(StateTransitionSteps::ExitSchedules))
-                    .add_systems(last_transition::<T>.pipe(run_transition::<T>).in_set(StateTransitionSteps::TransitionSchedules))
-                    .configure_sets(
-                        ApplyStateTransition::<T>::apply()
                         .in_set(StateTransitionSteps::DependentTransitions)
-                        $(.after(ApplyStateTransition::<$param::RawState>::apply()))*
+                        $(.after(ApplyStateTransition::<$param::RawState>::default()))*,
-                    );
+                    ExitSchedules::<T>::default()
+                        .in_set(StateTransitionSteps::ExitSchedules)
+                        $(.before(ExitSchedules::<$param::RawState>::default()))*,
+                    TransitionSchedules::<T>::default()
+                        .in_set(StateTransitionSteps::TransitionSchedules),
+                    EnterSchedules::<T>::default()
+                        .in_set(StateTransitionSteps::EnterSchedules)
+                        $(.after(EnterSchedules::<$param::RawState>::default()))*,
+                ));
+
+                schedule
+                    .add_systems(apply_state_transition.in_set(ApplyStateTransition::<T>::default()))
+                    .add_systems(last_transition::<T>.pipe(run_exit::<T>).in_set(ExitSchedules::<T>::default()))
+                    .add_systems(last_transition::<T>.pipe(run_transition::<T>).in_set(TransitionSchedules::<T>::default()))
+                    .add_systems(last_transition::<T>.pipe(run_enter::<T>).in_set(EnterSchedules::<T>::default()));
             }
         }
     };

crates/bevy_state/src/state/sub_states.rs

@@ -150,7 +150,7 @@ pub trait SubStates: States + FreelyMutableState {
     ///
     /// This can either be a single type that implements [`States`], or a tuple
     /// containing multiple types that implement [`States`], or any combination of
-    /// types implementing [`States`] and Options of types implementing [`States`]
+    /// types implementing [`States`] and Options of types implementing [`States`].
     type SourceStates: StateSet;
 
     /// This function gets called whenever one of the [`SourceStates`](Self::SourceStates) changes.

crates/bevy_state/src/state/transitions.rs

@@ -53,30 +53,58 @@ pub struct StateTransitionEvent<S: States> {
     pub entered: Option<S>,
 }
 
-/// Applies manual state transitions using [`NextState<S>`].
+/// Applies state transitions and runs transitions schedules in order.
 ///
 /// These system sets are run sequentially, in the order of the enum variants.
 #[derive(SystemSet, Clone, Debug, PartialEq, Eq, Hash)]
-pub enum StateTransitionSteps {
+pub(crate) enum StateTransitionSteps {
-    /// Parentless states apply their [`NextState<S>`].
+    /// States apply their transitions from [`NextState`] and compute functions based on their parent states.
-    RootTransitions,
-    /// States with parents apply their computation and [`NextState<S>`].
     DependentTransitions,
-    /// Exit schedules are executed.
+    /// Exit schedules are executed in leaf to root order
     ExitSchedules,
-    /// Transition schedules are executed.
+    /// Transition schedules are executed in arbitrary order.
     TransitionSchedules,
-    /// Enter schedules are executed.
+    /// Enter schedules are executed in root to leaf order.
     EnterSchedules,
 }
 
-/// Defines a system set to aid with dependent state ordering
 #[derive(SystemSet, Clone, Debug, PartialEq, Eq, Hash)]
-pub struct ApplyStateTransition<S: States>(PhantomData<S>);
+/// System set that runs exit schedule(s) for state `S`.
+pub struct ExitSchedules<S: States>(PhantomData<S>);
 
-impl<S: States> ApplyStateTransition<S> {
+impl<S: States> Default for ExitSchedules<S> {
-    pub(crate) fn apply() -> Self {
+    fn default() -> Self {
-        Self(PhantomData)
+        Self(Default::default())
+    }
+}
+
+#[derive(SystemSet, Clone, Debug, PartialEq, Eq, Hash)]
+/// System set that runs transition schedule(s) for state `S`.
+pub struct TransitionSchedules<S: States>(PhantomData<S>);
+
+impl<S: States> Default for TransitionSchedules<S> {
+    fn default() -> Self {
+        Self(Default::default())
+    }
+}
+
+#[derive(SystemSet, Clone, Debug, PartialEq, Eq, Hash)]
+/// System set that runs enter schedule(s) for state `S`.
+pub struct EnterSchedules<S: States>(PhantomData<S>);
+
+impl<S: States> Default for EnterSchedules<S> {
+    fn default() -> Self {
+        Self(Default::default())
+    }
+}
+
+/// System set that applies transitions for state `S`.
+#[derive(SystemSet, Clone, Debug, PartialEq, Eq, Hash)]
+pub(crate) struct ApplyStateTransition<S: States>(PhantomData<S>);
+
+impl<S: States> Default for ApplyStateTransition<S> {
+    fn default() -> Self {
+        Self(Default::default())
     }
 }
 
@@ -151,7 +179,6 @@ pub fn setup_state_transitions_in_world(
     let mut schedule = Schedule::new(StateTransition);
     schedule.configure_sets(
         (
-            StateTransitionSteps::RootTransitions,
             StateTransitionSteps::DependentTransitions,
             StateTransitionSteps::ExitSchedules,
             StateTransitionSteps::TransitionSchedules,

examples/state/custom_transitions.rs

@@ -13,10 +13,7 @@
 
 use std::marker::PhantomData;
 
-use bevy::{
+use bevy::{dev_tools::states::*, ecs::schedule::ScheduleLabel, prelude::*};
-    dev_tools::states::*, ecs::schedule::ScheduleLabel, prelude::*,
-    state::state::StateTransitionSteps,
-};
 
 use custom_transitions::*;
 
@@ -72,16 +69,16 @@ mod custom_transitions {
                     // 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`]
+                    // - [`ExitSchedules`] - Ran from leaf-states to root-states,
-                    // - [`StateTransitionSteps::TransitionSchedules`]
+                    // - [`TransitionSchedules`] - Ran in arbitrary order,
-                    // - [`StateTransitionSteps::EnterSchedules`]
+                    // - [`EnterSchedules`] - Ran from root-states to leaf-states.
-                    .in_set(StateTransitionSteps::EnterSchedules),
+                    .in_set(EnterSchedules::<S>::default()),
             )
             .add_systems(
                 StateTransition,
                 last_transition::<S>
                     .pipe(run_reexit::<S>)
-                    .in_set(StateTransitionSteps::ExitSchedules),
+                    .in_set(ExitSchedules::<S>::default()),
             );
         }
     }