Merge pull request bevyengine#88 from devil-ira/stage-yeet

Added `on_enter`, `on_update` and `on_exit` methods to `System(Set)Config`
tim-blackbird · Jan 23, 2023 · 4a87dde · 4a87dde
2 parents b6859c5 + 91e2b07
commit 4a87dde
Show file tree

Hide file tree

Showing 4 changed files with 127 additions and 29 deletions.
diff --git a/crates/bevy_app/src/app.rs b/crates/bevy_app/src/app.rs
@@ -188,7 +188,7 @@ impl App {
     /// Creates a new [`App`] with some default structure to enable core engine features.
     /// This is the preferred constructor for most use cases.
     ///
-    /// This calls [`App::add_default_schedules`] and [App::add_defaults_sets`].
+    /// This calls [`App::add_default_schedules`] and [`App::add_defaults_sets`].
     pub fn new() -> App {
         App::default()
     }
@@ -994,7 +994,7 @@ impl App {
 
 /// Methods for working with schedules
 impl App {
-    /// Adds a new `schedule` to the [`App`] under the provided `label.
+    /// Adds a new `schedule` to the [`App`] under the provided `label`.
     ///
     /// # Warning
     /// This method will overwrite any existing schedule at that label.
@@ -1006,7 +1006,7 @@ impl App {
         self
     }
 
-    /// Initializes a new empty `schedule` to the [`App`] under the provided `label if it does not exists.
+    /// Initializes a new empty `schedule` to the [`App`] under the provided `label` if it does not exists.
     ///
     /// See [`App::add_schedule`] to pass in a pre-constructed schedule.
     pub fn init_schedule(&mut self, label: impl ScheduleLabel) -> &mut Self {

diff --git a/crates/bevy_ecs/src/scheduling/config.rs b/crates/bevy_ecs/src/scheduling/config.rs
@@ -6,6 +6,7 @@ use crate::{
         condition::{BoxedCondition, Condition},
         graph_utils::{Ambiguity, Dependency, DependencyKind, GraphInfo},
         set::{BoxedSystemSet, IntoSystemSet, SystemSet},
+        state::{OnEnter, OnExit, OnUpdate, States},
     },
     system::{BoxedSystem, IntoSystem, System},
 };
@@ -102,6 +103,12 @@ pub trait IntoSystemSetConfig: sealed::IntoSystemSetConfig {
     /// The `Condition` will be evaluated at most once (per schedule run),
     /// the first time a system in this set prepares to run.
     fn run_if<P>(self, condition: impl Condition<P>) -> SystemSetConfig;
+    /// Add this set to the [`OnEnter(state)`](OnEnter) set.
+    fn on_enter(self, state: impl States) -> SystemSetConfig;
+    /// Add this set to the [`OnUpdate(state)`](OnUpdate) set.
+    fn on_update(self, state: impl States) -> SystemSetConfig;
+    /// Add this set to the [`OnExit(state)`](OnExit) set.
+    fn on_exit(self, state: impl States) -> SystemSetConfig;
     /// Suppress warnings and errors that would result from systems in this set having ambiguities
     /// (conflicting access but indeterminate order) with systems in `set`.
     fn ambiguous_with<M>(self, set: impl IntoSystemSet<M>) -> SystemSetConfig;
@@ -119,27 +126,39 @@ where
     }
 
     fn in_set(self, set: impl SystemSet) -> SystemSetConfig {
-        SystemSetConfig::new(Box::new(self)).in_set(set)
+        self.into_config().in_set(set)
     }
 
     fn before<M>(self, set: impl IntoSystemSet<M>) -> SystemSetConfig {
-        SystemSetConfig::new(Box::new(self)).before(set)
+        self.into_config().before(set)
     }
 
     fn after<M>(self, set: impl IntoSystemSet<M>) -> SystemSetConfig {
-        SystemSetConfig::new(Box::new(self)).after(set)
+        self.into_config().after(set)
     }
 
     fn run_if<P>(self, condition: impl Condition<P>) -> SystemSetConfig {
-        SystemSetConfig::new(Box::new(self)).run_if(condition)
+        self.into_config().run_if(condition)
+    }
+
+    fn on_enter(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_enter(state)
+    }
+
+    fn on_update(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_update(state)
+    }
+
+    fn on_exit(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_exit(state)
     }
 
     fn ambiguous_with<M>(self, set: impl IntoSystemSet<M>) -> SystemSetConfig {
-        SystemSetConfig::new(Box::new(self)).ambiguous_with(set)
+        self.into_config().ambiguous_with(set)
     }
 
     fn ambiguous_with_all(self) -> SystemSetConfig {
-        SystemSetConfig::new(Box::new(self)).ambiguous_with_all()
+        self.into_config().ambiguous_with_all()
     }
 }
 
@@ -149,27 +168,39 @@ impl IntoSystemSetConfig for BoxedSystemSet {
     }
 
     fn in_set(self, set: impl SystemSet) -> SystemSetConfig {
-        SystemSetConfig::new(self).in_set(set)
+        self.into_config().in_set(set)
     }
 
     fn before<M>(self, set: impl IntoSystemSet<M>) -> SystemSetConfig {
-        SystemSetConfig::new(self).before(set)
+        self.into_config().before(set)
     }
 
     fn after<M>(self, set: impl IntoSystemSet<M>) -> SystemSetConfig {
-        SystemSetConfig::new(self).after(set)
+        self.into_config().after(set)
     }
 
     fn run_if<P>(self, condition: impl Condition<P>) -> SystemSetConfig {
-        SystemSetConfig::new(self).run_if(condition)
+        self.into_config().run_if(condition)
+    }
+
+    fn on_enter(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_enter(state)
+    }
+
+    fn on_update(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_update(state)
+    }
+
+    fn on_exit(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_exit(state)
     }
 
     fn ambiguous_with<M>(self, set: impl IntoSystemSet<M>) -> SystemSetConfig {
-        SystemSetConfig::new(self).ambiguous_with(set)
+        self.into_config().ambiguous_with(set)
     }
 
     fn ambiguous_with_all(self) -> SystemSetConfig {
-        SystemSetConfig::new(self).ambiguous_with_all()
+        self.into_config().ambiguous_with_all()
     }
 }
 
@@ -208,6 +239,21 @@ impl IntoSystemSetConfig for SystemSetConfig {
         self
     }
 
+    fn on_enter(self, state: impl States) -> SystemSetConfig {
+        self.in_set(OnEnter(state));
+        self
+    }
+
+    fn on_update(self, state: impl States) -> SystemSetConfig {
+        self.in_set(OnUpdate(state));
+        self
+    }
+
+    fn on_exit(self, state: impl States) -> SystemSetConfig {
+        self.in_set(OnExit(state));
+        self
+    }
+
     fn ambiguous_with<M>(mut self, set: impl IntoSystemSet<M>) -> Self {
         ambiguous_with(&mut self.graph_info, Box::new(set.into_system_set()));
         self
@@ -238,6 +284,12 @@ pub trait IntoSystemConfig<Params>: sealed::IntoSystemConfig<Params> {
     /// The `Condition` will be evaluated at most once (per schedule run),
     /// when the system prepares to run.
     fn run_if<P>(self, condition: impl Condition<P>) -> SystemConfig;
+    /// Add this system to the [`OnEnter(state)`](OnEnter) set.
+    fn on_enter(self, state: impl States) -> SystemSetConfig;
+    /// Add this system to the [`OnUpdate(state)`](OnUpdate) set.
+    fn on_update(self, state: impl States) -> SystemSetConfig;
+    /// Add this system to the [`OnExit(state)`](OnExit) set.
+    fn on_exit(self, state: impl States) -> SystemSetConfig;
     /// Suppress warnings and errors that would result from this system having ambiguities
     /// (conflicting access but indeterminate order) with systems in `set`.
     fn ambiguous_with<M>(self, set: impl IntoSystemSet<M>) -> SystemConfig;
@@ -255,27 +307,39 @@ where
     }
 
     fn in_set(self, set: impl SystemSet) -> SystemConfig {
-        SystemConfig::new(Box::new(IntoSystem::into_system(self))).in_set(set)
+        self.into_config().in_set(set)
     }
 
     fn before<M>(self, set: impl IntoSystemSet<M>) -> SystemConfig {
-        SystemConfig::new(Box::new(IntoSystem::into_system(self))).before(set)
+        self.into_config().before(set)
     }
 
     fn after<M>(self, set: impl IntoSystemSet<M>) -> SystemConfig {
-        SystemConfig::new(Box::new(IntoSystem::into_system(self))).after(set)
+        self.into_config().after(set)
     }
 
     fn run_if<P>(self, condition: impl Condition<P>) -> SystemConfig {
-        SystemConfig::new(Box::new(IntoSystem::into_system(self))).run_if(condition)
+        self.into_config().run_if(condition)
+    }
+
+    fn on_enter(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_enter(state)
+    }
+
+    fn on_update(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_update(state)
+    }
+
+    fn on_exit(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_exit(state)
     }
 
     fn ambiguous_with<M>(self, set: impl IntoSystemSet<M>) -> SystemConfig {
-        SystemConfig::new(Box::new(IntoSystem::into_system(self))).ambiguous_with(set)
+        self.into_config().ambiguous_with(set)
     }
 
     fn ambiguous_with_all(self) -> SystemConfig {
-        SystemConfig::new(Box::new(IntoSystem::into_system(self))).ambiguous_with_all()
+        self.into_config().ambiguous_with_all()
     }
 }
 
@@ -285,27 +349,39 @@ impl IntoSystemConfig<()> for BoxedSystem<(), ()> {
     }
 
     fn in_set(self, set: impl SystemSet) -> SystemConfig {
-        SystemConfig::new(self).in_set(set)
+        self.into_config().in_set(set)
     }
 
     fn before<M>(self, set: impl IntoSystemSet<M>) -> SystemConfig {
-        SystemConfig::new(self).before(set)
+        self.into_config().before(set)
     }
 
     fn after<M>(self, set: impl IntoSystemSet<M>) -> SystemConfig {
-        SystemConfig::new(self).after(set)
+        self.into_config().after(set)
     }
 
     fn run_if<P>(self, condition: impl Condition<P>) -> SystemConfig {
-        SystemConfig::new(self).run_if(condition)
+        self.into_config().run_if(condition)
+    }
+
+    fn on_enter(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_enter(state)
+    }
+
+    fn on_update(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_update(state)
+    }
+
+    fn on_exit(self, state: impl States) -> SystemSetConfig {
+        self.into_config().on_exit(state)
     }
 
     fn ambiguous_with<M>(self, set: impl IntoSystemSet<M>) -> SystemConfig {
-        SystemConfig::new(self).ambiguous_with(set)
+        self.into_config().ambiguous_with(set)
     }
 
     fn ambiguous_with_all(self) -> SystemConfig {
-        SystemConfig::new(self).ambiguous_with_all()
+        self.into_config().ambiguous_with_all()
     }
 }
 
@@ -344,6 +420,21 @@ impl IntoSystemConfig<()> for SystemConfig {
         self
     }
 
+    fn on_enter(self, state: impl States) -> SystemSetConfig {
+        self.in_set(OnEnter(state));
+        self
+    }
+
+    fn on_update(self, state: impl States) -> SystemSetConfig {
+        self.in_set(OnUpdate(state));
+        self
+    }
+
+    fn on_exit(self, state: impl States) -> SystemSetConfig {
+        self.in_set(OnExit(state));
+        self
+    }
+
     fn ambiguous_with<M>(mut self, set: impl IntoSystemSet<M>) -> Self {
         ambiguous_with(&mut self.graph_info, Box::new(set.into_system_set()));
         self

diff --git a/crates/bevy_ecs/src/scheduling/state.rs b/crates/bevy_ecs/src/scheduling/state.rs
@@ -53,6 +53,13 @@ pub struct State<S: States>(pub S);
 #[derive(Resource, Default)]
 pub struct NextState<S: States>(pub Option<S>);
 
+impl<S: States> NextState<S> {
+    /// Queue the transition to a new `state`.
+    pub fn queue(&mut self, state: S) {
+        self.0 = Some(state);
+    }
+}
+
 /// If a new state is queued in [`NextState<S>`], this system:
 /// - Takes the new state value from [`NextState<S>`] and updates [`State<S>`].
 /// - Runs the [`OnExit(exited_state)`] schedule.

diff --git a/examples/ecs/state.rs b/examples/ecs/state.rs
@@ -86,7 +86,7 @@ fn setup_menu(mut commands: Commands, asset_server: Res<AssetServer>) {
 }
 
 fn menu(
-    mut state: ResMut<State<AppState>>,
+    mut next_state: ResMut<NextState<AppState>>,
     mut interaction_query: Query<
         (&Interaction, &mut BackgroundColor),
         (Changed<Interaction>, With<Button>),
@@ -96,7 +96,7 @@ fn menu(
         match *interaction {
             Interaction::Clicked => {
                 *color = PRESSED_BUTTON.into();
-                state.set(AppState::InGame).unwrap();
+                next_state.queue(AppState::InGame);
             }
             Interaction::Hovered => {
                 *color = HOVERED_BUTTON.into();