Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

allow using any stage as a startup stage #2248

Closed
wants to merge 1 commit into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
66 changes: 53 additions & 13 deletions crates/bevy_app/src/app_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -84,7 +84,9 @@ impl AppBuilder {
}

pub fn add_stage<S: Stage>(&mut self, label: impl StageLabel, stage: S) -> &mut Self {
let cloned_label = label.dyn_clone();
self.app.schedule.add_stage(label, stage);
self.add_startup_stage_boxed(cloned_label, SystemStage::parallel());
self
}

Expand All @@ -94,6 +96,7 @@ impl AppBuilder {
label: impl StageLabel,
stage: S,
) -> &mut Self {
self.add_startup_stage_after_boxed(&target, label.dyn_clone(), SystemStage::parallel());
self.app.schedule.add_stage_after(target, label, stage);
self
}
Expand All @@ -104,10 +107,51 @@ impl AppBuilder {
label: impl StageLabel,
stage: S,
) -> &mut Self {
self.add_startup_stage_before_boxed(&target, label.dyn_clone(), SystemStage::parallel());
self.app.schedule.add_stage_before(target, label, stage);
self
}

pub(crate) fn add_startup_stage_boxed<S: Stage>(
&mut self,
label: Box<dyn StageLabel>,
stage: S,
) -> &mut Self {
self.app
.schedule
.stage(CoreStage::Startup, |schedule: &mut Schedule| {
schedule.add_stage_boxed(label, stage)
});
self
}

pub(crate) fn add_startup_stage_after_boxed<S: Stage>(
&mut self,
target: &dyn StageLabel,
label: Box<dyn StageLabel>,
stage: S,
) -> &mut Self {
self.app
.schedule
.stage(CoreStage::Startup, |schedule: &mut Schedule| {
schedule.add_stage_after_boxed(target, label, stage)
});
self
}

pub(crate) fn add_startup_stage_before_boxed<S: Stage>(
&mut self,
target: &dyn StageLabel,
label: Box<dyn StageLabel>,
stage: S,
) -> &mut Self {
self.app
.schedule
.stage(CoreStage::Startup, |schedule: &mut Schedule| {
schedule.add_stage_before_boxed(target, label, stage)
});
self
}
pub fn add_startup_stage<S: Stage>(&mut self, label: impl StageLabel, stage: S) -> &mut Self {
self.app
.schedule
Expand Down Expand Up @@ -271,19 +315,15 @@ impl AppBuilder {
}

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())
self.add_stage(
CoreStage::Startup,
Schedule::default().with_run_criteria(RunOnce::default()),
)
.add_stage(CoreStage::First, 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())
}

/// Setup the application to manage events of type `T`.
Expand Down
45 changes: 43 additions & 2 deletions crates/bevy_app/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@ mod schedule_runner;
#[cfg(feature = "bevy_ci_testing")]
mod ci_testing;

use std::{any::Any, hash::Hash};

pub use app::*;
pub use app_builder::*;
pub use bevy_derive::DynamicPlugin;
Expand All @@ -23,7 +25,7 @@ pub mod prelude {
};
}

use bevy_ecs::schedule::StageLabel;
use bevy_ecs::schedule::{DynEq, DynHash, StageLabel};

/// The names of the default App stages
#[derive(Debug, Hash, PartialEq, Eq, Clone, StageLabel)]
Expand All @@ -42,8 +44,9 @@ pub enum CoreStage {
/// Name of app stage that runs after all other app stages
Last,
}

/// The names of the default App startup stages
#[derive(Debug, Hash, PartialEq, Eq, Clone, StageLabel)]
#[derive(Debug, PartialEq, Clone)]
pub enum StartupStage {
/// Name of app stage that runs once before the startup stage
PreStartup,
Expand All @@ -52,3 +55,41 @@ pub enum StartupStage {
/// Name of app stage that runs once after the startup stage
PostStartup,
}

impl StageLabel for StartupStage {
fn dyn_clone(&self) -> Box<dyn StageLabel> {
Box::new(Clone::clone(self))
}
}

impl DynEq for StartupStage {
fn as_any(&self) -> &dyn Any {
self
}

fn dyn_eq(&self, other: &dyn DynEq) -> bool {
match self {
StartupStage::PreStartup => CoreStage::PreUpdate.dyn_eq(other),
StartupStage::Startup => CoreStage::Update.dyn_eq(other),
StartupStage::PostStartup => CoreStage::PostUpdate.dyn_eq(other),
}
}
}
impl DynHash for StartupStage {
fn as_dyn_eq(&self) -> &dyn bevy_ecs::schedule::DynEq {
match self {
StartupStage::PreStartup => &CoreStage::PreUpdate,
StartupStage::Startup => &CoreStage::Update,
StartupStage::PostStartup => &CoreStage::PostUpdate,
}
}

fn dyn_hash(&self, mut state: &mut dyn std::hash::Hasher) {
match self {
StartupStage::PreStartup => CoreStage::PreUpdate.hash(&mut state),
StartupStage::Startup => CoreStage::Update.hash(&mut state),
StartupStage::PostStartup => CoreStage::PostUpdate.hash(&mut state),
}
CoreStage::Update.type_id().hash(&mut state);
}
}
29 changes: 27 additions & 2 deletions crates/bevy_ecs/src/schedule/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -85,6 +85,11 @@ impl Schedule {

pub fn add_stage<S: Stage>(&mut self, label: impl StageLabel, stage: S) -> &mut Self {
let label: Box<dyn StageLabel> = Box::new(label);
self.add_stage_boxed(label, stage)
}

#[doc(hidden)]
pub fn add_stage_boxed<S: Stage>(&mut self, label: Box<dyn StageLabel>, stage: S) -> &mut Self {
self.stage_order.push(label.clone());
let prev = self.stages.insert(label.clone(), Box::new(stage));
if prev.is_some() {
Expand All @@ -101,11 +106,21 @@ impl Schedule {
) -> &mut Self {
let label: Box<dyn StageLabel> = Box::new(label);
let target = &target as &dyn StageLabel;
self.add_stage_after_boxed(target, label, stage)
}

#[doc(hidden)]
pub fn add_stage_after_boxed<S: Stage>(
&mut self,
target: &dyn StageLabel,
label: Box<dyn StageLabel>,
stage: S,
) -> &mut Self {
let target_index = self
.stage_order
.iter()
.enumerate()
.find(|(_i, stage_label)| &***stage_label == target)
.find(|(_i, stage_label)| target == &***stage_label)
.map(|(i, _)| i)
.unwrap_or_else(|| panic!("Target stage does not exist: {:?}.", target));

Expand All @@ -125,11 +140,21 @@ impl Schedule {
) -> &mut Self {
let label: Box<dyn StageLabel> = Box::new(label);
let target = &target as &dyn StageLabel;
self.add_stage_before_boxed(target, label, stage)
}

#[doc(hidden)]
pub fn add_stage_before_boxed<S: Stage>(
&mut self,
target: &dyn StageLabel,
label: Box<dyn StageLabel>,
stage: S,
) -> &mut Self {
let target_index = self
.stage_order
.iter()
.enumerate()
.find(|(_i, stage_label)| &***stage_label == target)
.find(|(_i, stage_label)| target == &***stage_label)
.map(|(i, _)| i)
.unwrap_or_else(|| panic!("Target stage does not exist: {:?}.", target));

Expand Down