From 868e8b4530a96db678e1a5b2316038e6f3c3330b Mon Sep 17 00:00:00 2001
From: Christian Hughes <9044780+ItsDoot@users.noreply.github.com>
Date: Fri, 10 Mar 2023 03:06:23 -0600
Subject: [PATCH] Add OnTransition schedule that is ran between OnExit and
 OnEnter (#7936)

---
 crates/bevy_ecs/src/lib.rs            |  4 ++--
 crates/bevy_ecs/src/schedule/state.rs | 25 +++++++++++++++++++++++--
 2 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs
index fce9800306f6d..1c574befe7065 100644
--- a/crates/bevy_ecs/src/lib.rs
+++ b/crates/bevy_ecs/src/lib.rs
@@ -40,8 +40,8 @@ pub mod prelude {
         schedule::{
             apply_state_transition, apply_system_buffers, common_conditions::*, Condition,
             IntoSystemConfig, IntoSystemConfigs, IntoSystemSet, IntoSystemSetConfig,
-            IntoSystemSetConfigs, NextState, OnEnter, OnExit, OnUpdate, Schedule, Schedules, State,
-            States, SystemSet,
+            IntoSystemSetConfigs, NextState, OnEnter, OnExit, OnTransition, OnUpdate, Schedule,
+            Schedules, State, States, SystemSet,
         },
         system::{
             adapter as system_adapter,
diff --git a/crates/bevy_ecs/src/schedule/state.rs b/crates/bevy_ecs/src/schedule/state.rs
index 92f7c3864ce85..abb001e90a8d0 100644
--- a/crates/bevy_ecs/src/schedule/state.rs
+++ b/crates/bevy_ecs/src/schedule/state.rs
@@ -4,7 +4,7 @@ use std::mem;
 
 use crate as bevy_ecs;
 use crate::change_detection::DetectChangesMut;
-use crate::schedule::{ScheduleLabel, SystemSet};
+use crate::schedule::{ScheduleLabel, Schedules, SystemSet};
 use crate::system::Resource;
 use crate::world::World;
 
@@ -54,6 +54,18 @@ pub struct OnEnter<S: States>(pub S);
 #[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
 pub struct OnExit<S: States>(pub S);
 
+/// The label of a [`Schedule`](super::Schedule) that **only** runs whenever [`State<S>`]
+/// exits the `from` state, AND enters the `to` state.
+///
+/// Systems added to this schedule are always ran *after* [`OnExit`], and *before* [`OnEnter`].
+#[derive(ScheduleLabel, Clone, Debug, PartialEq, Eq, Hash)]
+pub struct OnTransition<S: States> {
+    /// The state being exited.
+    pub from: S,
+    /// The state being entered.
+    pub to: S,
+}
+
 /// A [`SystemSet`] that will run within `CoreSet::Update` when this state is active.
 ///
 /// This set, when created via `App::add_state`, is configured with both a base set and a run condition.
@@ -105,7 +117,16 @@ pub fn apply_state_transition<S: States>(world: &mut World) {
         next_state_resource.set_changed();
 
         let exited = mem::replace(&mut world.resource_mut::<State<S>>().0, entered.clone());
-        world.run_schedule(OnExit(exited));
+        world.run_schedule(OnExit(exited.clone()));
+
+        let transition_schedule = OnTransition {
+            from: exited,
+            to: entered.clone(),
+        };
+        if world.resource::<Schedules>().contains(&transition_schedule) {
+            world.run_schedule(transition_schedule);
+        }
+
         world.run_schedule(OnEnter(entered));
     }
 }