Remove the Stage trait

crates/bevy_app/src/app.rs

@@ -2,9 +2,7 @@ use crate::{CoreStage, Events, Plugin, PluginGroup, PluginGroupBuilder, StartupS
 use bevy_ecs::{
     component::{Component, ComponentDescriptor},
     prelude::{FromWorld, IntoExclusiveSystem},
-    schedule::{
-        IntoSystemDescriptor, RunOnce, Schedule, Stage, StageLabel, State, SystemSet, SystemStage,
-    },
+    schedule::{IntoSystemDescriptor, Schedule, StageLabel, State, SystemSet, SystemStage},
     world::World,
 };
 use bevy_utils::tracing::debug;
@@ -40,6 +38,7 @@ pub struct App {
     pub world: World,
     pub runner: Box<dyn Fn(App)>,
     pub schedule: Schedule,
+    pub startup_schedule: Schedule,
 }
 
 impl Default for App {
@@ -49,6 +48,7 @@ impl Default for App {
         app.init_resource::<bevy_reflect::TypeRegistryArc>();
 
         app.add_default_stages()
+            .add_default_startup_stages()
             .add_event::<AppExit>()
             .add_system_to_stage(CoreStage::Last, World::clear_trackers.exclusive_system());
 
@@ -71,6 +71,7 @@ impl App {
             world: Default::default(),
             schedule: Default::default(),
             runner: Box::new(run_once),
+            startup_schedule: Default::default(),
         }
     }
 
@@ -109,69 +110,60 @@ impl App {
         (runner)(app);
     }
 
-    pub fn add_stage<S: Stage>(&mut self, label: impl StageLabel, stage: S) -> &mut Self {
+    pub fn add_stage(&mut self, label: impl StageLabel, stage: SystemStage) -> &mut Self {
         self.schedule.add_stage(label, stage);
         self
     }
 
-    pub fn add_stage_after<S: Stage>(
+    pub fn add_stage_after(
         &mut self,
         target: impl StageLabel,
         label: impl StageLabel,
-        stage: S,
+        stage: SystemStage,
     ) -> &mut Self {
         self.schedule.add_stage_after(target, label, stage);
         self
     }
 
-    pub fn add_stage_before<S: Stage>(
+    pub fn add_stage_before(
         &mut self,
         target: impl StageLabel,
         label: impl StageLabel,
-        stage: S,
+        stage: SystemStage,
     ) -> &mut Self {
         self.schedule.add_stage_before(target, label, stage);
         self
     }
 
-    pub fn add_startup_stage<S: Stage>(&mut self, label: impl StageLabel, stage: S) -> &mut Self {
-        self.schedule
-            .stage(CoreStage::Startup, |schedule: &mut Schedule| {
-                schedule.add_stage(label, stage)
-            });
+    pub fn add_startup_stage(&mut self, label: impl StageLabel, stage: SystemStage) -> &mut Self {
+        self.startup_schedule.add_stage(label, stage);
         self
     }
 
-    pub fn add_startup_stage_after<S: Stage>(
+    pub fn add_startup_stage_after(
         &mut self,
         target: impl StageLabel,
         label: impl StageLabel,
-        stage: S,
+        stage: SystemStage,
     ) -> &mut Self {
-        self.schedule
-            .stage(CoreStage::Startup, |schedule: &mut Schedule| {
-                schedule.add_stage_after(target, label, stage)
-            });
+        self.startup_schedule.add_stage_after(target, label, stage);
         self
     }
 
-    pub fn add_startup_stage_before<S: Stage>(
+    pub fn add_startup_stage_before(
         &mut self,
         target: impl StageLabel,
         label: impl StageLabel,
-        stage: S,
+        stage: SystemStage,
     ) -> &mut Self {
-        self.schedule
-            .stage(CoreStage::Startup, |schedule: &mut Schedule| {
-                schedule.add_stage_before(target, label, stage)
-            });
+        self.startup_schedule.add_stage_before(target, label, stage);
         self
     }
 
-    pub fn stage<T: Stage, F: FnOnce(&mut T) -> &mut T>(
+    pub fn stage(
         &mut self,
         label: impl StageLabel,
-        func: F,
+        func: impl FnOnce(&mut SystemStage) -> &mut SystemStage,
     ) -> &mut Self {
         self.schedule.stage(label, func);
         self
@@ -266,10 +258,8 @@ impl App {
         stage_label: impl StageLabel,
         system: impl IntoSystemDescriptor<Params>,
     ) -> &mut Self {
-        self.schedule
-            .stage(CoreStage::Startup, |schedule: &mut Schedule| {
-                schedule.add_system_to_stage(stage_label, system)
-            });
+        self.startup_schedule
+            .add_system_to_stage(stage_label, system);
         self
     }
 
@@ -278,10 +268,8 @@ impl App {
         stage_label: impl StageLabel,
         system_set: SystemSet,
     ) -> &mut Self {
-        self.schedule
-            .stage(CoreStage::Startup, |schedule: &mut Schedule| {
-                schedule.add_system_set_to_stage(stage_label, system_set)
-            });
+        self.startup_schedule
+            .add_system_set_to_stage(stage_label, system_set);
         self
     }
 
@@ -312,20 +300,18 @@ impl App {
 
     pub fn add_default_stages(&mut self) -> &mut Self {
         self.add_stage(CoreStage::First, SystemStage::parallel())
-            .add_stage(
-                CoreStage::Startup,
-                Schedule::default()
-                    .with_run_criteria(RunOnce::default())
-                    .with_stage(StartupStage::PreStartup, SystemStage::parallel())
-                    .with_stage(StartupStage::Startup, SystemStage::parallel())
-                    .with_stage(StartupStage::PostStartup, SystemStage::parallel()),
-            )
             .add_stage(CoreStage::PreUpdate, SystemStage::parallel())
             .add_stage(CoreStage::Update, SystemStage::parallel())
             .add_stage(CoreStage::PostUpdate, SystemStage::parallel())
             .add_stage(CoreStage::Last, SystemStage::parallel())
     }
 
+    pub fn add_default_startup_stages(&mut self) -> &mut Self {
+        self.add_startup_stage(StartupStage::PreStartup, SystemStage::parallel())
+            .add_startup_stage(StartupStage::Startup, SystemStage::parallel())
+            .add_startup_stage(StartupStage::PostStartup, SystemStage::parallel())
+    }
+
     /// Setup the application to manage events of type `T`.
     ///
     /// This is done by adding a `Resource` of type `Events::<T>`,
@@ -579,6 +565,7 @@ impl App {
 }
 
 fn run_once(mut app: App) {
+    app.startup_schedule.run(&mut app.world);
     app.update();
 }
 

crates/bevy_app/src/schedule_runner.rs

@@ -58,6 +58,7 @@ impl Plugin for ScheduleRunnerPlugin {
             .to_owned();
         app.set_runner(move |mut app: App| {
             let mut app_exit_event_reader = ManualEventReader::<AppExit>::default();
+            app.startup_schedule.run(&mut app.world);
             match settings.run_mode {
                 RunMode::Once => {
                     app.update();

crates/bevy_core/src/label.rs

@@ -122,7 +122,7 @@ pub(crate) fn entity_labels_system(
 #[cfg(test)]
 mod tests {
     use bevy_ecs::{
-        schedule::{Schedule, Stage, SystemStage},
+        schedule::{Schedule, SystemStage},
         world::World,
     };
 

crates/bevy_ecs/src/lib.rs

@@ -26,7 +26,7 @@ pub mod prelude {
         schedule::{
             AmbiguitySetLabel, ExclusiveSystemDescriptorCoercion, ParallelSystemDescriptorCoercion,
             RunCriteria, RunCriteriaDescriptorCoercion, RunCriteriaLabel, RunCriteriaPiping,
-            Schedule, Stage, StageLabel, State, SystemLabel, SystemSet, SystemStage,
+            Schedule, StageLabel, State, SystemLabel, SystemSet, SystemStage,
         },
         system::{
             Commands, ConfigurableSystem, In, IntoChainSystem, IntoExclusiveSystem, IntoSystem,

crates/bevy_ecs/src/schedule/executor_parallel.rs

@@ -315,7 +315,7 @@ enum SchedulingEvent {
 mod tests {
     use super::SchedulingEvent::{self, *};
     use crate::{
-        schedule::{SingleThreadedExecutor, Stage, SystemStage},
+        schedule::{SingleThreadedExecutor, SystemStage},
         system::{NonSend, Query, Res, ResMut},
         world::World,
     };

crates/bevy_ecs/src/schedule/mod.rs

@@ -27,32 +27,32 @@ use bevy_utils::HashMap;
 
 #[derive(Default)]
 pub struct Schedule {
-    stages: HashMap<BoxedStageLabel, Box<dyn Stage>>,
+    stages: HashMap<BoxedStageLabel, SystemStage>,
     stage_order: Vec<BoxedStageLabel>,
     run_criteria: BoxedRunCriteria,
 }
 
 impl Schedule {
-    pub fn with_stage<S: Stage>(mut self, label: impl StageLabel, stage: S) -> Self {
+    pub fn with_stage(mut self, label: impl StageLabel, stage: SystemStage) -> Self {
         self.add_stage(label, stage);
         self
     }
 
-    pub fn with_stage_after<S: Stage>(
+    pub fn with_stage_after(
         mut self,
         target: impl StageLabel,
         label: impl StageLabel,
-        stage: S,
+        stage: SystemStage,
     ) -> Self {
         self.add_stage_after(target, label, stage);
         self
     }
 
-    pub fn with_stage_before<S: Stage>(
+    pub fn with_stage_before(
         mut self,
         target: impl StageLabel,
         label: impl StageLabel,
-        stage: S,
+        stage: SystemStage,
     ) -> Self {
         self.add_stage_before(target, label, stage);
         self
@@ -80,21 +80,21 @@ impl Schedule {
         self
     }
 
-    pub fn add_stage<S: Stage>(&mut self, label: impl StageLabel, stage: S) -> &mut Self {
+    pub fn add_stage(&mut self, label: impl StageLabel, stage: SystemStage) -> &mut Self {
         let label: Box<dyn StageLabel> = Box::new(label);
         self.stage_order.push(label.clone());
-        let prev = self.stages.insert(label.clone(), Box::new(stage));
+        let prev = self.stages.insert(label.clone(), stage);
         if prev.is_some() {
             panic!("Stage already exists: {:?}.", label);
         }
         self
     }
 
-    pub fn add_stage_after<S: Stage>(
+    pub fn add_stage_after(
         &mut self,
         target: impl StageLabel,
         label: impl StageLabel,
-        stage: S,
+        stage: SystemStage,
     ) -> &mut Self {
         let label: Box<dyn StageLabel> = Box::new(label);
         let target = &target as &dyn StageLabel;
@@ -107,18 +107,18 @@ impl Schedule {
             .unwrap_or_else(|| panic!("Target stage does not exist: {:?}.", target));
 
         self.stage_order.insert(target_index + 1, label.clone());
-        let prev = self.stages.insert(label.clone(), Box::new(stage));
+        let prev = self.stages.insert(label.clone(), stage);
         if prev.is_some() {
             panic!("Stage already exists: {:?}.", label);
         }
         self
     }
 
-    pub fn add_stage_before<S: Stage>(
+    pub fn add_stage_before(
         &mut self,
         target: impl StageLabel,
         label: impl StageLabel,
-        stage: S,
+        stage: SystemStage,
     ) -> &mut Self {
         let label: Box<dyn StageLabel> = Box::new(label);
         let target = &target as &dyn StageLabel;
@@ -131,7 +131,7 @@ impl Schedule {
             .unwrap_or_else(|| panic!("Target stage does not exist: {:?}.", target));
 
         self.stage_order.insert(target_index, label.clone());
-        let prev = self.stages.insert(label.clone(), Box::new(stage));
+        let prev = self.stages.insert(label.clone(), stage);
         if prev.is_some() {
             panic!("Stage already exists: {:?}.", label);
         }
@@ -154,7 +154,7 @@ impl Schedule {
         }
 
         let stage = self
-            .get_stage_mut::<SystemStage>(&stage_label)
+            .get_stage_mut(&stage_label)
             .unwrap_or_else(move || stage_not_found(&stage_label));
         stage.add_system(system);
         self
@@ -170,31 +170,27 @@ impl Schedule {
         })
     }
 
-    pub fn stage<T: Stage, F: FnOnce(&mut T) -> &mut T>(
+    pub fn stage(
         &mut self,
         label: impl StageLabel,
-        func: F,
+        func: impl FnOnce(&mut SystemStage) -> &mut SystemStage,
     ) -> &mut Self {
-        let stage = self.get_stage_mut::<T>(&label).unwrap_or_else(move || {
+        let stage = self.get_stage_mut(&label).unwrap_or_else(move || {
             panic!("stage '{:?}' does not exist or is the wrong type", label)
         });
         func(stage);
         self
     }
 
-    pub fn get_stage<T: Stage>(&self, label: &dyn StageLabel) -> Option<&T> {
-        self.stages
-            .get(label)
-            .and_then(|stage| stage.downcast_ref::<T>())
+    pub fn get_stage(&self, label: &dyn StageLabel) -> Option<&SystemStage> {
+        self.stages.get(label)
     }
 
-    pub fn get_stage_mut<T: Stage>(&mut self, label: &dyn StageLabel) -> Option<&mut T> {
-        self.stages
-            .get_mut(label)
-            .and_then(|stage| stage.downcast_mut::<T>())
+    pub fn get_stage_mut(&mut self, label: &dyn StageLabel) -> Option<&mut SystemStage> {
+        self.stages.get_mut(label)
     }
 
-    pub fn run_once(&mut self, world: &mut World) {
+    pub fn run(&mut self, world: &mut World) {
         for label in self.stage_order.iter() {
             #[cfg(feature = "trace")]
             let stage_span =
@@ -207,29 +203,9 @@ impl Schedule {
     }
 
     /// Iterates over all of schedule's stages and their labels, in execution order.
-    pub fn iter_stages(&self) -> impl Iterator<Item = (&dyn StageLabel, &dyn Stage)> {
+    pub fn iter_stages(&self) -> impl Iterator<Item = (&dyn StageLabel, &SystemStage)> {
         self.stage_order
             .iter()
-            .map(move |label| (&**label, &*self.stages[label]))
-    }
-}
-
-impl Stage for Schedule {
-    fn run(&mut self, world: &mut World) {
-        loop {
-            match self.run_criteria.should_run(world) {
-                ShouldRun::No => return,
-                ShouldRun::Yes => {
-                    self.run_once(world);
-                    return;
-                }
-                ShouldRun::YesAndCheckAgain => {
-                    self.run_once(world);
-                }
-                ShouldRun::NoAndCheckAgain => {
-                    panic!("`NoAndCheckAgain` would loop infinitely in this situation.")
-                }
-            }
-        }
+            .map(move |label| (&**label, &self.stages[label]))
     }
 }