Improve usability of StateStage and cut down on "magic" (#1059)

Improve usability of StateStage and cut down on "magic"

crates/bevy_app/src/app_builder.rs

@@ -1,5 +1,3 @@
-use std::any::Any;
-
 use crate::{
     app::{App, AppExit},
     event::Events,
@@ -7,8 +5,8 @@ use crate::{
     stage, startup_stage, PluginGroup, PluginGroupBuilder,
 };
 use bevy_ecs::{
-    clear_trackers_system, FromResources, IntoStage, IntoSystem, Resource, Resources, RunOnce,
+    clear_trackers_system, FromResources, IntoSystem, Resource, Resources, RunOnce, Schedule,
-    Schedule, Stage, State, StateStage, System, SystemStage, World,
+    Stage, StateStage, System, SystemStage, World,
 };
 use bevy_utils::tracing::debug;
 
@@ -56,16 +54,12 @@ impl AppBuilder {
         self
     }
 
-    pub fn add_stage<Params, S: IntoStage<Params>>(
+    pub fn add_stage<S: Stage>(&mut self, name: &'static str, stage: S) -> &mut Self {
-        &mut self,
-        name: &'static str,
-        stage: S,
-    ) -> &mut Self {
         self.app.schedule.add_stage(name, stage);
         self
     }
 
-    pub fn add_stage_after<Params, S: IntoStage<Params>>(
+    pub fn add_stage_after<S: Stage>(
         &mut self,
         target: &'static str,
         name: &'static str,
@@ -75,7 +69,7 @@ impl AppBuilder {
         self
     }
 
-    pub fn add_stage_before<Params, S: IntoStage<Params>>(
+    pub fn add_stage_before<S: Stage>(
         &mut self,
         target: &'static str,
         name: &'static str,
@@ -85,11 +79,7 @@ impl AppBuilder {
         self
     }
 
-    pub fn add_startup_stage<Params, S: IntoStage<Params>>(
+    pub fn add_startup_stage<S: Stage>(&mut self, name: &'static str, stage: S) -> &mut Self {
-        &mut self,
-        name: &'static str,
-        stage: S,
-    ) -> &mut Self {
         self.app
             .schedule
             .stage(stage::STARTUP, |schedule: &mut Schedule| {
@@ -98,7 +88,7 @@ impl AppBuilder {
         self
     }
 
-    pub fn add_startup_stage_after<Params, S: IntoStage<Params>>(
+    pub fn add_startup_stage_after<S: Stage>(
         &mut self,
         target: &'static str,
         name: &'static str,
@@ -112,7 +102,7 @@ impl AppBuilder {
         self
     }
 
-    pub fn add_startup_stage_before<Params, S: IntoStage<Params>>(
+    pub fn add_startup_stage_before<S: Stage>(
         &mut self,
         target: &'static str,
         name: &'static str,
@@ -143,6 +133,51 @@ impl AppBuilder {
         self.add_system_to_stage(stage::UPDATE, system)
     }
 
+    pub fn on_state_enter<T: Clone + Resource, S, Params, IntoS>(
+        &mut self,
+        stage: &str,
+        state: T,
+        system: IntoS,
+    ) -> &mut Self
+    where
+        S: System<In = (), Out = ()>,
+        IntoS: IntoSystem<Params, S>,
+    {
+        self.stage(stage, |stage: &mut StateStage<T>| {
+            stage.on_state_enter(state, system)
+        })
+    }
+
+    pub fn on_state_update<T: Clone + Resource, S, Params, IntoS>(
+        &mut self,
+        stage: &str,
+        state: T,
+        system: IntoS,
+    ) -> &mut Self
+    where
+        S: System<In = (), Out = ()>,
+        IntoS: IntoSystem<Params, S>,
+    {
+        self.stage(stage, |stage: &mut StateStage<T>| {
+            stage.on_state_update(state, system)
+        })
+    }
+
+    pub fn on_state_exit<T: Clone + Resource, S, Params, IntoS>(
+        &mut self,
+        stage: &str,
+        state: T,
+        system: IntoS,
+    ) -> &mut Self
+    where
+        S: System<In = (), Out = ()>,
+        IntoS: IntoSystem<Params, S>,
+    {
+        self.stage(stage, |stage: &mut StateStage<T>| {
+            stage.on_state_exit(state, system)
+        })
+    }
+
     pub fn add_startup_system_to_stage<S, Params, IntoS>(
         &mut self,
         stage_name: &'static str,
@@ -207,53 +242,6 @@ impl AppBuilder {
             .add_system_to_stage(stage::EVENT, Events::<T>::update_system)
     }
 
-    pub fn state_stage_name<T: Any>() -> String {
-        format!("state({})", std::any::type_name::<T>())
-    }
-
-    pub fn add_state<T: Clone + Resource>(&mut self, initial: T) -> &mut Self {
-        self.add_resource(State::new(initial));
-        self.app.schedule.add_stage_after(
-            stage::UPDATE,
-            &Self::state_stage_name::<T>(),
-            StateStage::<T>::default(),
-        );
-        self
-    }
-
-    pub fn on_state_enter<T: Clone + Resource, Params, S: IntoStage<Params>>(
-        &mut self,
-        value: T,
-        stage: S,
-    ) -> &mut Self {
-        self.stage(
-            &Self::state_stage_name::<T>(),
-            |state_stage: &mut StateStage<T>| state_stage.on_state_enter(value, stage),
-        )
-    }
-
-    pub fn on_state_update<T: Clone + Resource, Params, S: IntoStage<Params>>(
-        &mut self,
-        value: T,
-        stage: S,
-    ) -> &mut Self {
-        self.stage(
-            &Self::state_stage_name::<T>(),
-            |state_stage: &mut StateStage<T>| state_stage.on_state_update(value, stage),
-        )
-    }
-
-    pub fn on_state_exit<T: Clone + Resource, Params, S: IntoStage<Params>>(
-        &mut self,
-        value: T,
-        stage: S,
-    ) -> &mut Self {
-        self.stage(
-            &Self::state_stage_name::<T>(),
-            |state_stage: &mut StateStage<T>| state_stage.on_state_exit(value, stage),
-        )
-    }
-
     /// Adds a resource to the current [App] and overwrites any resource previously added of the same type.
     pub fn add_resource<T>(&mut self, resource: T) -> &mut Self
     where

crates/bevy_ecs/src/lib.rs

@@ -14,7 +14,7 @@ pub mod prelude {
     pub use crate::{
         core::WorldBuilderSource,
         resource::{ChangedRes, FromResources, Local, Res, ResMut, Resource, Resources},
-        schedule::{Schedule, State, SystemStage},
+        schedule::{Schedule, State, StateStage, SystemStage},
         system::{Commands, IntoSystem, Query, System},
         Added, Bundle, Changed, Component, Entity, In, IntoChainSystem, Mut, Mutated, Or, QuerySet,
         Ref, RefMut, With, Without, World,

crates/bevy_ecs/src/schedule/mod.rs

@@ -18,27 +18,17 @@ pub struct Schedule {
 }
 
 impl Schedule {
-    pub fn with_stage<Params, S: IntoStage<Params>>(mut self, name: &str, stage: S) -> Self {
+    pub fn with_stage<S: Stage>(mut self, name: &str, stage: S) -> Self {
-        self.add_stage(name, stage.into_stage());
+        self.add_stage(name, stage);
         self
     }
 
-    pub fn with_stage_after<Params, S: IntoStage<Params>>(
+    pub fn with_stage_after<S: Stage>(mut self, target: &str, name: &str, stage: S) -> Self {
-        mut self,
-        target: &str,
-        name: &str,
-        stage: S,
-    ) -> Self {
         self.add_stage_after(target, name, stage);
         self
     }
 
-    pub fn with_stage_before<Params, S: IntoStage<Params>>(
+    pub fn with_stage_before<S: Stage>(mut self, target: &str, name: &str, stage: S) -> Self {
-        mut self,
-        target: &str,
-        name: &str,
-        stage: S,
-    ) -> Self {
         self.add_stage_before(target, name, stage);
         self
     }
@@ -75,19 +65,13 @@ impl Schedule {
         self
     }
 
-    pub fn add_stage<Params, S: IntoStage<Params>>(&mut self, name: &str, stage: S) -> &mut Self {
+    pub fn add_stage<S: Stage>(&mut self, name: &str, stage: S) -> &mut Self {
         self.stage_order.push(name.to_string());
-        self.stages
+        self.stages.insert(name.to_string(), Box::new(stage));
-            .insert(name.to_string(), Box::new(stage.into_stage()));
         self
     }
 
-    pub fn add_stage_after<Params, S: IntoStage<Params>>(
+    pub fn add_stage_after<S: Stage>(&mut self, target: &str, name: &str, stage: S) -> &mut Self {
-        &mut self,
-        target: &str,
-        name: &str,
-        stage: S,
-    ) -> &mut Self {
         if self.stages.get(name).is_some() {
             panic!("Stage already exists: {}.", name);
         }
@@ -100,18 +84,12 @@ impl Schedule {
             .map(|(i, _)| i)
             .unwrap_or_else(|| panic!("Target stage does not exist: {}.", target));
 
-        self.stages
+        self.stages.insert(name.to_string(), Box::new(stage));
-            .insert(name.to_string(), Box::new(stage.into_stage()));
         self.stage_order.insert(target_index + 1, name.to_string());
         self
     }
 
-    pub fn add_stage_before<Params, S: IntoStage<Params>>(
+    pub fn add_stage_before<S: Stage>(&mut self, target: &str, name: &str, stage: S) -> &mut Self {
-        &mut self,
-        target: &str,
-        name: &str,
-        stage: S,
-    ) -> &mut Self {
         if self.stages.get(name).is_some() {
             panic!("Stage already exists: {}.", name);
         }
@@ -124,8 +102,7 @@ impl Schedule {
             .map(|(i, _)| i)
             .unwrap_or_else(|| panic!("Target stage does not exist: {}.", target));
 
-        self.stages
+        self.stages.insert(name.to_string(), Box::new(stage));
-            .insert(name.to_string(), Box::new(stage.into_stage()));
         self.stage_order.insert(target_index, name.to_string());
         self
     }

crates/bevy_ecs/src/schedule/stage.rs

@@ -174,29 +174,6 @@ impl<S: System<In = (), Out = ()>> From<S> for SystemStage {
     }
 }
 
-pub trait IntoStage<Params> {
-    type Stage: Stage;
-    fn into_stage(self) -> Self::Stage;
-}
-
-impl<Params, S: System<In = (), Out = ()>, IntoS: IntoSystem<Params, S>> IntoStage<(Params, S)>
-    for IntoS
-{
-    type Stage = SystemStage;
-
-    fn into_stage(self) -> Self::Stage {
-        SystemStage::single(self)
-    }
-}
-
-impl<S: Stage> IntoStage<()> for S {
-    type Stage = S;
-
-    fn into_stage(self) -> Self::Stage {
-        self
-    }
-}
-
 pub struct RunOnce {
     ran: bool,
     system_id: SystemId,

crates/bevy_ecs/src/schedule/state.rs

@@ -1,13 +1,22 @@
-use crate::{IntoStage, Resource, Resources, Stage, World};
+use crate::{IntoSystem, Resource, Resources, Stage, System, SystemStage, World};
 use bevy_utils::HashMap;
 use std::{mem::Discriminant, ops::Deref};
 use thiserror::Error;
 
-#[derive(Default)]
 pub(crate) struct StateStages {
-    update: Option<Box<dyn Stage>>,
+    update: Box<dyn Stage>,
-    enter: Option<Box<dyn Stage>>,
+    enter: Box<dyn Stage>,
-    exit: Option<Box<dyn Stage>>,
+    exit: Box<dyn Stage>,
+}
+
+impl Default for StateStages {
+    fn default() -> Self {
+        Self {
+            enter: Box::new(SystemStage::parallel()),
+            update: Box::new(SystemStage::parallel()),
+            exit: Box::new(SystemStage::parallel()),
+        }
+    }
 }
 
 pub struct StateStage<T> {
@@ -24,76 +33,113 @@ impl<T> Default for StateStage<T> {
 
 #[allow(clippy::mem_discriminant_non_enum)]
 impl<T> StateStage<T> {
-    pub fn with_on_state_enter<Params, S: IntoStage<Params>>(mut self, state: T, stage: S) -> Self {
+    pub fn set_enter_stage<S: Stage>(&mut self, state: T, stage: S) -> &mut Self {
-        self.on_state_enter(state, stage);
+        let stages = self.state_stages(state);
+        stages.enter = Box::new(stage);
         self
     }
 
-    pub fn with_on_state_exit<Params, S: IntoStage<Params>>(mut self, state: T, stage: S) -> Self {
+    pub fn set_exit_stage<S: Stage>(&mut self, state: T, stage: S) -> &mut Self {
-        self.on_state_exit(state, stage);
+        let stages = self.state_stages(state);
+        stages.exit = Box::new(stage);
         self
     }
 
-    pub fn with_on_state_update<Params, S: IntoStage<Params>>(
+    pub fn set_update_stage<S: Stage>(&mut self, state: T, stage: S) -> &mut Self {
-        mut self,
+        let stages = self.state_stages(state);
-        state: T,
+        stages.update = Box::new(stage);
-        stage: S,
-    ) -> Self {
-        self.on_state_update(state, stage);
         self
     }
 
-    pub fn on_state_enter<Params, S: IntoStage<Params>>(
+    pub fn on_state_enter<Params, S: System<In = (), Out = ()>, IntoS: IntoSystem<Params, S>>(
         &mut self,
         state: T,
-        stage: S,
+        system: IntoS,
     ) -> &mut Self {
-        let stages = self
+        self.enter_stage(state, |system_stage: &mut SystemStage| {
-            .stages
+            system_stage.add_system(system)
-            .entry(std::mem::discriminant(&state))
+        })
-            .or_default();
-        stages.enter = Some(Box::new(stage.into_stage()));
-        self
     }
 
-    pub fn on_state_exit<Params, S: IntoStage<Params>>(&mut self, state: T, stage: S) -> &mut Self {
+    pub fn on_state_exit<Params, S: System<In = (), Out = ()>, IntoS: IntoSystem<Params, S>>(
-        let stages = self
-            .stages
-            .entry(std::mem::discriminant(&state))
-            .or_default();
-        stages.exit = Some(Box::new(stage.into_stage()));
-        self
-    }
-
-    pub fn on_state_update<Params, S: IntoStage<Params>>(
         &mut self,
         state: T,
-        stage: S,
+        system: IntoS,
     ) -> &mut Self {
-        let stages = self
+        self.exit_stage(state, |system_stage: &mut SystemStage| {
-            .stages
+            system_stage.add_system(system)
-            .entry(std::mem::discriminant(&state))
+        })
-            .or_default();
+    }
-        stages.update = Some(Box::new(stage.into_stage()));
+
+    pub fn on_state_update<Params, S: System<In = (), Out = ()>, IntoS: IntoSystem<Params, S>>(
+        &mut self,
+        state: T,
+        system: IntoS,
+    ) -> &mut Self {
+        self.update_stage(state, |system_stage: &mut SystemStage| {
+            system_stage.add_system(system)
+        })
+    }
+
+    pub fn enter_stage<S: Stage, F: FnOnce(&mut S) -> &mut S>(
+        &mut self,
+        state: T,
+        func: F,
+    ) -> &mut Self {
+        let stages = self.state_stages(state);
+        func(
+            stages
+                .enter
+                .downcast_mut()
+                .expect("'Enter' stage does not match the given type"),
+        );
         self
     }
+
+    pub fn exit_stage<S: Stage, F: FnOnce(&mut S) -> &mut S>(
+        &mut self,
+        state: T,
+        func: F,
+    ) -> &mut Self {
+        let stages = self.state_stages(state);
+        func(
+            stages
+                .exit
+                .downcast_mut()
+                .expect("'Exit' stage does not match the given type"),
+        );
+        self
+    }
+
+    pub fn update_stage<S: Stage, F: FnOnce(&mut S) -> &mut S>(
+        &mut self,
+        state: T,
+        func: F,
+    ) -> &mut Self {
+        let stages = self.state_stages(state);
+        func(
+            stages
+                .update
+                .downcast_mut()
+                .expect("'Update' stage does not match the given type"),
+        );
+        self
+    }
+
+    fn state_stages(&mut self, state: T) -> &mut StateStages {
+        self.stages
+            .entry(std::mem::discriminant(&state))
+            .or_default()
+    }
 }
 
 #[allow(clippy::mem_discriminant_non_enum)]
 impl<T: Resource + Clone> Stage for StateStage<T> {
     fn initialize(&mut self, world: &mut World, resources: &mut Resources) {
         for state_stages in self.stages.values_mut() {
-            if let Some(ref mut enter) = state_stages.enter {
+            state_stages.enter.initialize(world, resources);
-                enter.initialize(world, resources);
+            state_stages.update.initialize(world, resources);
-            }
+            state_stages.exit.initialize(world, resources);
-
-            if let Some(ref mut update) = state_stages.update {
-                update.initialize(world, resources);
-            }
-
-            if let Some(ref mut exit) = state_stages.exit {
-                exit.initialize(world, resources);
-            }
         }
     }
 
@@ -116,33 +162,21 @@ impl<T: Resource + Clone> Stage for StateStage<T> {
             // if next_stage is Some, we just applied a new state
             if let Some(next_stage) = next_stage {
                 if next_stage != current_stage {
-                    if let Some(exit_current) = self
+                    if let Some(current_state_stages) = self.stages.get_mut(&current_stage) {
-                        .stages
+                        current_state_stages.exit.run(world, resources);
-                        .get_mut(&current_stage)
-                        .and_then(|stage| stage.exit.as_mut())
-                    {
-                        exit_current.run(world, resources);
                     }
                 }
 
-                if let Some(enter_next) = self
+                if let Some(next_state_stages) = self.stages.get_mut(&next_stage) {
-                    .stages
+                    next_state_stages.enter.run(world, resources);
-                    .get_mut(&next_stage)
-                    .and_then(|stage| stage.enter.as_mut())
-                {
-                    enter_next.run(world, resources);
                 }
             } else {
                 break current_stage;
             }
         };
 
-        if let Some(update_current) = self
+        if let Some(current_state_stages) = self.stages.get_mut(&current_stage) {
-            .stages
+            current_state_stages.update.run(world, resources);
-            .get_mut(&current_stage)
-            .and_then(|stage| stage.update.as_mut())
-        {
-            update_current.run(world, resources);
         }
     }
 }

crates/bevy_ecs/src/system/into_system.rs

@@ -330,7 +330,7 @@ mod tests {
         let mut update = SystemStage::parallel();
         update.add_system(incr_e_on_flip);
         schedule.add_stage("update", update);
-        schedule.add_stage("clear_trackers", clear_trackers_system);
+        schedule.add_stage("clear_trackers", SystemStage::single(clear_trackers_system));
 
         schedule.initialize_and_run(&mut world, &mut resources);
         assert_eq!(*(world.get::<i32>(ent).unwrap()), 1);
@@ -364,7 +364,7 @@ mod tests {
         let mut update = SystemStage::parallel();
         update.add_system(incr_e_on_flip);
         schedule.add_stage("update", update);
-        schedule.add_stage("clear_trackers", clear_trackers_system);
+        schedule.add_stage("clear_trackers", SystemStage::single(clear_trackers_system));
 
         schedule.initialize_and_run(&mut world, &mut resources);
         assert_eq!(*(world.get::<i32>(ent).unwrap()), 1);

examples/2d/texture_atlas.rs

@@ -5,13 +5,16 @@ fn main() {
     App::build()
         .init_resource::<RpgSpriteHandles>()
         .add_plugins(DefaultPlugins)
-        .add_state(AppState::Setup)
+        .add_resource(State::new(AppState::Setup))
-        .on_state_enter(AppState::Setup, load_textures)
+        .add_stage_after(stage::UPDATE, STAGE, StateStage::<AppState>::default())
-        .on_state_update(AppState::Setup, check_textures)
+        .on_state_enter(STAGE, AppState::Setup, load_textures)
-        .on_state_enter(AppState::Finshed, setup)
+        .on_state_update(STAGE, AppState::Setup, check_textures)
+        .on_state_enter(STAGE, AppState::Finshed, setup)
         .run();
 }
 
+const STAGE: &str = "app_state";
+
 #[derive(Clone)]
 enum AppState {
     Setup,

examples/ecs/state.rs

@@ -5,20 +5,19 @@ fn main() {
     App::build()
         .add_plugins(DefaultPlugins)
         .init_resource::<ButtonMaterials>()
-        .add_state(AppState::Menu)
+        .add_resource(State::new(AppState::Menu))
-        .on_state_enter(AppState::Menu, setup_menu)
+        .add_stage_after(stage::UPDATE, STAGE, StateStage::<AppState>::default())
-        .on_state_update(AppState::Menu, menu)
+        .on_state_enter(STAGE, AppState::Menu, setup_menu)
-        .on_state_exit(AppState::Menu, cleanup_menu)
+        .on_state_update(STAGE, AppState::Menu, menu)
-        .on_state_enter(AppState::InGame, setup_game)
+        .on_state_exit(STAGE, AppState::Menu, cleanup_menu)
-        .on_state_update(
+        .on_state_enter(STAGE, AppState::InGame, setup_game)
-            AppState::InGame,
+        .on_state_update(STAGE, AppState::InGame, movement)
-            SystemStage::parallel()
+        .on_state_update(STAGE, AppState::InGame, change_color)
-                .with_system(movement)
-                .with_system(change_color),
-        )
         .run();
 }
 
+const STAGE: &str = "app_state";
+
 #[derive(Clone)]
 enum AppState {
     Menu,