From b12e3bf3bbec14d2ea987377ea2bb2acc20a4db3 Mon Sep 17 00:00:00 2001 From: Carter Anderson Date: Mon, 14 Dec 2020 17:13:22 -0800 Subject: [PATCH] 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 | 118 +++++++-------- crates/bevy_ecs/src/lib.rs | 2 +- crates/bevy_ecs/src/schedule/mod.rs | 43 ++---- crates/bevy_ecs/src/schedule/stage.rs | 23 --- crates/bevy_ecs/src/schedule/state.rs | 170 +++++++++++++--------- crates/bevy_ecs/src/system/into_system.rs | 4 +- examples/2d/texture_atlas.rs | 11 +- examples/ecs/state.rs | 21 ++- 8 files changed, 185 insertions(+), 207 deletions(-) diff --git a/crates/bevy_app/src/app_builder.rs b/crates/bevy_app/src/app_builder.rs index a25198bcab..3c0b620ba4 100644 --- a/crates/bevy_app/src/app_builder.rs +++ b/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, - Schedule, Stage, State, StateStage, System, SystemStage, World, + clear_trackers_system, FromResources, IntoSystem, Resource, Resources, RunOnce, Schedule, + Stage, StateStage, System, SystemStage, World, }; use bevy_utils::tracing::debug; @@ -56,16 +54,12 @@ impl AppBuilder { self } - pub fn add_stage>( - &mut self, - name: &'static str, - stage: S, - ) -> &mut Self { + pub fn add_stage(&mut self, name: &'static str, stage: S) -> &mut Self { self.app.schedule.add_stage(name, stage); self } - pub fn add_stage_after>( + pub fn add_stage_after( &mut self, target: &'static str, name: &'static str, @@ -75,7 +69,7 @@ impl AppBuilder { self } - pub fn add_stage_before>( + pub fn add_stage_before( &mut self, target: &'static str, name: &'static str, @@ -85,11 +79,7 @@ impl AppBuilder { self } - pub fn add_startup_stage>( - &mut self, - name: &'static str, - stage: S, - ) -> &mut Self { + pub fn add_startup_stage(&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>( + pub fn add_startup_stage_after( &mut self, target: &'static str, name: &'static str, @@ -112,7 +102,7 @@ impl AppBuilder { self } - pub fn add_startup_stage_before>( + pub fn add_startup_stage_before( &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( + &mut self, + stage: &str, + state: T, + system: IntoS, + ) -> &mut Self + where + S: System, + IntoS: IntoSystem, + { + self.stage(stage, |stage: &mut StateStage| { + stage.on_state_enter(state, system) + }) + } + + pub fn on_state_update( + &mut self, + stage: &str, + state: T, + system: IntoS, + ) -> &mut Self + where + S: System, + IntoS: IntoSystem, + { + self.stage(stage, |stage: &mut StateStage| { + stage.on_state_update(state, system) + }) + } + + pub fn on_state_exit( + &mut self, + stage: &str, + state: T, + system: IntoS, + ) -> &mut Self + where + S: System, + IntoS: IntoSystem, + { + self.stage(stage, |stage: &mut StateStage| { + stage.on_state_exit(state, system) + }) + } + pub fn add_startup_system_to_stage( &mut self, stage_name: &'static str, @@ -207,53 +242,6 @@ impl AppBuilder { .add_system_to_stage(stage::EVENT, Events::::update_system) } - pub fn state_stage_name() -> String { - format!("state({})", std::any::type_name::()) - } - - pub fn add_state(&mut self, initial: T) -> &mut Self { - self.add_resource(State::new(initial)); - self.app.schedule.add_stage_after( - stage::UPDATE, - &Self::state_stage_name::(), - StateStage::::default(), - ); - self - } - - pub fn on_state_enter>( - &mut self, - value: T, - stage: S, - ) -> &mut Self { - self.stage( - &Self::state_stage_name::(), - |state_stage: &mut StateStage| state_stage.on_state_enter(value, stage), - ) - } - - pub fn on_state_update>( - &mut self, - value: T, - stage: S, - ) -> &mut Self { - self.stage( - &Self::state_stage_name::(), - |state_stage: &mut StateStage| state_stage.on_state_update(value, stage), - ) - } - - pub fn on_state_exit>( - &mut self, - value: T, - stage: S, - ) -> &mut Self { - self.stage( - &Self::state_stage_name::(), - |state_stage: &mut StateStage| 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(&mut self, resource: T) -> &mut Self where diff --git a/crates/bevy_ecs/src/lib.rs b/crates/bevy_ecs/src/lib.rs index cdf8dfa5d6..0931d9f64d 100644 --- a/crates/bevy_ecs/src/lib.rs +++ b/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, diff --git a/crates/bevy_ecs/src/schedule/mod.rs b/crates/bevy_ecs/src/schedule/mod.rs index 4b05dabd18..30e2acd20d 100644 --- a/crates/bevy_ecs/src/schedule/mod.rs +++ b/crates/bevy_ecs/src/schedule/mod.rs @@ -18,27 +18,17 @@ pub struct Schedule { } impl Schedule { - pub fn with_stage>(mut self, name: &str, stage: S) -> Self { - self.add_stage(name, stage.into_stage()); + pub fn with_stage(mut self, name: &str, stage: S) -> Self { + self.add_stage(name, stage); self } - pub fn with_stage_after>( - mut self, - target: &str, - name: &str, - stage: S, - ) -> Self { + pub fn with_stage_after(mut self, target: &str, name: &str, stage: S) -> Self { self.add_stage_after(target, name, stage); self } - pub fn with_stage_before>( - mut self, - target: &str, - name: &str, - stage: S, - ) -> Self { + pub fn with_stage_before(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>(&mut self, name: &str, stage: S) -> &mut Self { + pub fn add_stage(&mut self, name: &str, stage: S) -> &mut Self { self.stage_order.push(name.to_string()); - self.stages - .insert(name.to_string(), Box::new(stage.into_stage())); + self.stages.insert(name.to_string(), Box::new(stage)); self } - pub fn add_stage_after>( - &mut self, - target: &str, - name: &str, - stage: S, - ) -> &mut Self { + pub fn add_stage_after(&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 - .insert(name.to_string(), Box::new(stage.into_stage())); + self.stages.insert(name.to_string(), Box::new(stage)); self.stage_order.insert(target_index + 1, name.to_string()); self } - pub fn add_stage_before>( - &mut self, - target: &str, - name: &str, - stage: S, - ) -> &mut Self { + pub fn add_stage_before(&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 - .insert(name.to_string(), Box::new(stage.into_stage())); + self.stages.insert(name.to_string(), Box::new(stage)); self.stage_order.insert(target_index, name.to_string()); self } diff --git a/crates/bevy_ecs/src/schedule/stage.rs b/crates/bevy_ecs/src/schedule/stage.rs index f15d60ebee..4917ac67f8 100644 --- a/crates/bevy_ecs/src/schedule/stage.rs +++ b/crates/bevy_ecs/src/schedule/stage.rs @@ -174,29 +174,6 @@ impl> From for SystemStage { } } -pub trait IntoStage { - type Stage: Stage; - fn into_stage(self) -> Self::Stage; -} - -impl, IntoS: IntoSystem> IntoStage<(Params, S)> - for IntoS -{ - type Stage = SystemStage; - - fn into_stage(self) -> Self::Stage { - SystemStage::single(self) - } -} - -impl IntoStage<()> for S { - type Stage = S; - - fn into_stage(self) -> Self::Stage { - self - } -} - pub struct RunOnce { ran: bool, system_id: SystemId, diff --git a/crates/bevy_ecs/src/schedule/state.rs b/crates/bevy_ecs/src/schedule/state.rs index d6c6963032..c023dc6625 100644 --- a/crates/bevy_ecs/src/schedule/state.rs +++ b/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>, - enter: Option>, - exit: Option>, + update: Box, + enter: Box, + exit: Box, +} + +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 { @@ -24,76 +33,113 @@ impl Default for StateStage { #[allow(clippy::mem_discriminant_non_enum)] impl StateStage { - pub fn with_on_state_enter>(mut self, state: T, stage: S) -> Self { - self.on_state_enter(state, stage); + pub fn set_enter_stage(&mut self, state: T, stage: S) -> &mut Self { + let stages = self.state_stages(state); + stages.enter = Box::new(stage); self } - pub fn with_on_state_exit>(mut self, state: T, stage: S) -> Self { - self.on_state_exit(state, stage); + pub fn set_exit_stage(&mut self, state: T, stage: S) -> &mut Self { + let stages = self.state_stages(state); + stages.exit = Box::new(stage); self } - pub fn with_on_state_update>( - mut self, - state: T, - stage: S, - ) -> Self { - self.on_state_update(state, stage); + pub fn set_update_stage(&mut self, state: T, stage: S) -> &mut Self { + let stages = self.state_stages(state); + stages.update = Box::new(stage); self } - pub fn on_state_enter>( + pub fn on_state_enter, IntoS: IntoSystem>( &mut self, state: T, - stage: S, + system: IntoS, ) -> &mut Self { - let stages = self - .stages - .entry(std::mem::discriminant(&state)) - .or_default(); - stages.enter = Some(Box::new(stage.into_stage())); - self + self.enter_stage(state, |system_stage: &mut SystemStage| { + system_stage.add_system(system) + }) } - pub fn on_state_exit>(&mut self, state: T, stage: S) -> &mut Self { - 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>( + pub fn on_state_exit, IntoS: IntoSystem>( &mut self, state: T, - stage: S, + system: IntoS, ) -> &mut Self { - let stages = self - .stages - .entry(std::mem::discriminant(&state)) - .or_default(); - stages.update = Some(Box::new(stage.into_stage())); + self.exit_stage(state, |system_stage: &mut SystemStage| { + system_stage.add_system(system) + }) + } + + pub fn on_state_update, IntoS: IntoSystem>( + &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 &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 &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 &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 Stage for StateStage { 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 { - enter.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); - } + state_stages.enter.initialize(world, resources); + state_stages.update.initialize(world, resources); + state_stages.exit.initialize(world, resources); } } @@ -116,33 +162,21 @@ impl Stage for StateStage { // 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 - .stages - .get_mut(¤t_stage) - .and_then(|stage| stage.exit.as_mut()) - { - exit_current.run(world, resources); + if let Some(current_state_stages) = self.stages.get_mut(¤t_stage) { + current_state_stages.exit.run(world, resources); } } - if let Some(enter_next) = self - .stages - .get_mut(&next_stage) - .and_then(|stage| stage.enter.as_mut()) - { - enter_next.run(world, resources); + if let Some(next_state_stages) = self.stages.get_mut(&next_stage) { + next_state_stages.enter.run(world, resources); } } else { break current_stage; } }; - if let Some(update_current) = self - .stages - .get_mut(¤t_stage) - .and_then(|stage| stage.update.as_mut()) - { - update_current.run(world, resources); + if let Some(current_state_stages) = self.stages.get_mut(¤t_stage) { + current_state_stages.update.run(world, resources); } } } diff --git a/crates/bevy_ecs/src/system/into_system.rs b/crates/bevy_ecs/src/system/into_system.rs index 967fe9666d..364232cedb 100644 --- a/crates/bevy_ecs/src/system/into_system.rs +++ b/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::(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::(ent).unwrap()), 1); diff --git a/examples/2d/texture_atlas.rs b/examples/2d/texture_atlas.rs index 08823067d4..3ee454a820 100644 --- a/examples/2d/texture_atlas.rs +++ b/examples/2d/texture_atlas.rs @@ -5,13 +5,16 @@ fn main() { App::build() .init_resource::() .add_plugins(DefaultPlugins) - .add_state(AppState::Setup) - .on_state_enter(AppState::Setup, load_textures) - .on_state_update(AppState::Setup, check_textures) - .on_state_enter(AppState::Finshed, setup) + .add_resource(State::new(AppState::Setup)) + .add_stage_after(stage::UPDATE, STAGE, StateStage::::default()) + .on_state_enter(STAGE, AppState::Setup, load_textures) + .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, diff --git a/examples/ecs/state.rs b/examples/ecs/state.rs index 07ae083bb1..1b79193c4f 100644 --- a/examples/ecs/state.rs +++ b/examples/ecs/state.rs @@ -5,20 +5,19 @@ fn main() { App::build() .add_plugins(DefaultPlugins) .init_resource::() - .add_state(AppState::Menu) - .on_state_enter(AppState::Menu, setup_menu) - .on_state_update(AppState::Menu, menu) - .on_state_exit(AppState::Menu, cleanup_menu) - .on_state_enter(AppState::InGame, setup_game) - .on_state_update( - AppState::InGame, - SystemStage::parallel() - .with_system(movement) - .with_system(change_color), - ) + .add_resource(State::new(AppState::Menu)) + .add_stage_after(stage::UPDATE, STAGE, StateStage::::default()) + .on_state_enter(STAGE, AppState::Menu, setup_menu) + .on_state_update(STAGE, AppState::Menu, menu) + .on_state_exit(STAGE, AppState::Menu, cleanup_menu) + .on_state_enter(STAGE, AppState::InGame, setup_game) + .on_state_update(STAGE, AppState::InGame, movement) + .on_state_update(STAGE, AppState::InGame, change_color) .run(); } +const STAGE: &str = "app_state"; + #[derive(Clone)] enum AppState { Menu,