From 0727610a340e8d64c4836901910b05804927c2b0 Mon Sep 17 00:00:00 2001 From: Mike Hsu Date: Tue, 27 Dec 2022 22:52:44 -0800 Subject: [PATCH] remove extract stage to run during extract --- crates/bevy_app/src/app.rs | 8 ++++ crates/bevy_app/src/plugin.rs | 9 +++++ crates/bevy_ecs/src/schedule/mod.rs | 11 ++++++ crates/bevy_render/src/lib.rs | 58 ++++++++++++++++++----------- 4 files changed, 65 insertions(+), 21 deletions(-) diff --git a/crates/bevy_app/src/app.rs b/crates/bevy_app/src/app.rs index d6b903bd2d433d..99dc0762842ad2 100644 --- a/crates/bevy_app/src/app.rs +++ b/crates/bevy_app/src/app.rs @@ -189,6 +189,14 @@ impl App { if app.is_building_plugin { panic!("App::run() was called from within Plugin::Build(), which is not allowed."); } + + // temporarily remove the plugin registry to run each plugin's setup function on app. + let mut plugin_registry = std::mem::take(&mut app.plugin_registry); + for plugin in &plugin_registry { + plugin.setup(&mut app); + } + std::mem::swap(&mut app.plugin_registry, &mut plugin_registry); + let runner = std::mem::replace(&mut app.runner, Box::new(run_once)); (runner)(app); } diff --git a/crates/bevy_app/src/plugin.rs b/crates/bevy_app/src/plugin.rs index 7e6e0c575a4d7d..abaf24c0339d28 100644 --- a/crates/bevy_app/src/plugin.rs +++ b/crates/bevy_app/src/plugin.rs @@ -16,10 +16,19 @@ use std::any::Any; pub trait Plugin: Downcast + Any + Send + Sync { /// Configures the [`App`] to which this plugin is added. fn build(&self, app: &mut App); + + /// Runs after all plugins are built, but before the app runner is called. + /// This can be useful if you have some resource that other plugins need during their build step, + /// but after build you want to remove it and send it to another thread. + fn setup(&self, _app: &mut App) { + // do nothing + } + /// Configures a name for the [`Plugin`] which is primarily used for debugging. fn name(&self) -> &str { std::any::type_name::() } + /// If the plugin can be meaningfully instantiated several times in an [`App`](crate::App), /// override this method to return `false`. fn is_unique(&self) -> bool { diff --git a/crates/bevy_ecs/src/schedule/mod.rs b/crates/bevy_ecs/src/schedule/mod.rs index 222339606fec34..1a00453ef9f987 100644 --- a/crates/bevy_ecs/src/schedule/mod.rs +++ b/crates/bevy_ecs/src/schedule/mod.rs @@ -361,6 +361,17 @@ impl Schedule { .and_then(|stage| stage.downcast_mut::()) } + /// Remove a [`Stage`] from the schedule + pub fn remove_stage(&mut self, stage_label: impl StageLabel) -> Option> { + let label = stage_label.as_label(); + + let Some(index) = self.stage_order.iter().position(|x| *x == label) else { + return None; + }; + self.stage_order.remove(index); + self.stages.remove(&label) + } + /// Executes each [`Stage`] contained in the schedule, one at a time. pub fn run_once(&mut self, world: &mut World) { for label in &self.stage_order { diff --git a/crates/bevy_render/src/lib.rs b/crates/bevy_render/src/lib.rs index 635217273dc6c9..c74ab11915447c 100644 --- a/crates/bevy_render/src/lib.rs +++ b/crates/bevy_render/src/lib.rs @@ -92,6 +92,10 @@ pub enum RenderStage { Cleanup, } +/// Resource for holding the extract stage of the rendering schedule +#[derive(Resource)] +pub struct ExtractStage(pub SystemStage); + /// The simulation [`World`] of the application, stored as a resource. /// This resource is only available during [`RenderStage::Extract`] and not /// during command application of that stage. @@ -266,6 +270,20 @@ impl Plugin for RenderPlugin { .register_type::() .register_type::(); } + + fn setup(&self, app: &mut App) { + if let Ok(render_app) = app.get_sub_app_mut(RenderApp) { + // move the extract stage to a resource so render_app.run() does not run it. + let stage = render_app + .schedule + .remove_stage(RenderStage::Extract) + .unwrap() + .downcast::() + .unwrap(); + + render_app.world.insert_resource(ExtractStage(*stage)); + } + } } /// A "scratch" world used to avoid allocating new worlds every frame when @@ -276,25 +294,23 @@ struct ScratchMainWorld(World); /// Executes the [`Extract`](RenderStage::Extract) stage of the renderer. /// This updates the render world with the extracted ECS data of the current frame. fn extract(app_world: &mut World, render_app: &mut App) { - let extract = render_app - .schedule - .get_stage_mut::(RenderStage::Extract) - .unwrap(); - - // temporarily add the app world to the render world as a resource - let scratch_world = app_world.remove_resource::().unwrap(); - let inserted_world = std::mem::replace(app_world, scratch_world.0); - let running_world = &mut render_app.world; - running_world.insert_resource(MainWorld(inserted_world)); - - extract.run(running_world); - // move the app world back, as if nothing happened. - let inserted_world = running_world.remove_resource::().unwrap(); - let scratch_world = std::mem::replace(app_world, inserted_world.0); - app_world.insert_resource(ScratchMainWorld(scratch_world)); - - // Note: We apply buffers (read, Commands) after the `MainWorld` has been removed from the render app's world - // so that in future, pipelining will be able to do this too without any code relying on it. - // see - extract.apply_buffers(running_world); + render_app + .world + .resource_scope(|render_world, mut extract_stage: Mut| { + // temporarily add the app world to the render world as a resource + let scratch_world = app_world.remove_resource::().unwrap(); + let inserted_world = std::mem::replace(app_world, scratch_world.0); + render_world.insert_resource(MainWorld(inserted_world)); + + extract_stage.0.run(render_world); + // move the app world back, as if nothing happened. + let inserted_world = render_world.remove_resource::().unwrap(); + let scratch_world = std::mem::replace(app_world, inserted_world.0); + app_world.insert_resource(ScratchMainWorld(scratch_world)); + + // Note: We apply buffers (read, Commands) after the `MainWorld` has been removed from the render app's world + // so that in future, pipelining will be able to do this too without any code relying on it. + // see + extract_stage.0.apply_buffers(render_world); + }); }