ExtractComponentPlugin: Support extracting the component only for vis…

…ible entities
bevyengine · May 2, 2022 · 9b3d816 · 9b3d816
1 parent 1794083
commit 9b3d816
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 21 deletions.
diff --git a/crates/bevy_pbr/src/material.rs b/crates/bevy_pbr/src/material.rs
@@ -10,14 +10,14 @@ use bevy_ecs::{
     prelude::World,
     system::{
         lifetimeless::{Read, SQuery, SRes},
-        Commands, Local, Query, Res, ResMut, SystemParamItem,
+        Query, Res, ResMut, SystemParamItem,
     },
     world::FromWorld,
 };
 use bevy_render::{
     mesh::{Mesh, MeshVertexBufferLayout},
-    prelude::ComputedVisibility,
     render_asset::{RenderAsset, RenderAssetPlugin, RenderAssets},
+    render_component::ExtractComponentPlugin,
     render_phase::{
         AddRenderCommand, DrawFunctions, EntityRenderCommand, RenderCommandResult, RenderPhase,
         SetItemPipeline, TrackedRenderPass,
@@ -204,6 +204,7 @@ impl<M: SpecializedMaterial> Default for MaterialPlugin<M> {
 impl<M: SpecializedMaterial> Plugin for MaterialPlugin<M> {
     fn build(&self, app: &mut App) {
         app.add_asset::<M>()
+            .add_plugin(ExtractComponentPlugin::<Handle<M>>::extract_visible())
             .add_plugin(RenderAssetPlugin::<M>::default());
         if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
             render_app
@@ -212,27 +213,11 @@ impl<M: SpecializedMaterial> Plugin for MaterialPlugin<M> {
                 .add_render_command::<AlphaMask3d, DrawMaterial<M>>()
                 .init_resource::<MaterialPipeline<M>>()
                 .init_resource::<SpecializedMeshPipelines<MaterialPipeline<M>>>()
-                .add_system_to_stage(RenderStage::Extract, extract_materials::<M>)
                 .add_system_to_stage(RenderStage::Queue, queue_material_meshes::<M>);
         }
     }
 }
 
-fn extract_materials<M: SpecializedMaterial>(
-    mut commands: Commands,
-    query: Query<(Entity, &ComputedVisibility, &Handle<M>)>,
-    mut prev_len: Local<usize>,
-) {
-    let mut materials = Vec::with_capacity(*prev_len);
-    for (entity, computed_visibility, material) in query.iter() {
-        if computed_visibility.is_visible {
-            materials.push((entity, (material.clone_weak(),)));
-        }
-    }
-    *prev_len = materials.len();
-    commands.insert_or_spawn_batch(materials);
-}
-
 #[derive(Eq, PartialEq, Clone, Hash)]
 pub struct MaterialPipelineKey<T> {
     pub mesh_key: MeshPipelineKey,

diff --git a/crates/bevy_render/src/render_component.rs b/crates/bevy_render/src/render_component.rs
@@ -1,6 +1,7 @@
 use crate::{
     render_resource::{std140::AsStd140, DynamicUniformVec},
     renderer::{RenderDevice, RenderQueue},
+    view::ComputedVisibility,
     RenderApp, RenderStage,
 };
 use bevy_app::{App, Plugin};
@@ -131,18 +132,38 @@ fn prepare_uniform_components<C: Component>(
 ///
 /// Therefore it sets up the [`RenderStage::Extract`](crate::RenderStage::Extract) step
 /// for the specified [`ExtractComponent`].
-pub struct ExtractComponentPlugin<C, F = ()>(PhantomData<fn() -> (C, F)>);
+pub struct ExtractComponentPlugin<C, F = ()> {
+    only_extract_visible: bool,
+    marker: PhantomData<fn() -> (C, F)>,
+}
 
 impl<C, F> Default for ExtractComponentPlugin<C, F> {
     fn default() -> Self {
-        Self(PhantomData)
+        Self {
+            only_extract_visible: false,
+            marker: PhantomData,
+        }
+    }
+}
+
+impl<C, F> ExtractComponentPlugin<C, F> {
+    pub fn extract_visible() -> Self {
+        Self {
+            only_extract_visible: true,
+            marker: PhantomData,
+        }
     }
 }
 
 impl<C: ExtractComponent> Plugin for ExtractComponentPlugin<C> {
     fn build(&self, app: &mut App) {
         if let Ok(render_app) = app.get_sub_app_mut(RenderApp) {
-            render_app.add_system_to_stage(RenderStage::Extract, extract_components::<C>);
+            if self.only_extract_visible {
+                render_app
+                    .add_system_to_stage(RenderStage::Extract, extract_visible_components::<C>);
+            } else {
+                render_app.add_system_to_stage(RenderStage::Extract, extract_components::<C>);
+            }
         }
     }
 }
@@ -170,3 +191,19 @@ fn extract_components<C: ExtractComponent>(
     *previous_len = values.len();
     commands.insert_or_spawn_batch(values);
 }
+
+/// This system extracts all visible components of the corresponding [`ExtractComponent`] type.
+fn extract_visible_components<C: ExtractComponent>(
+    mut commands: Commands,
+    mut previous_len: Local<usize>,
+    mut query: StaticSystemParam<Query<(Entity, Read<ComputedVisibility>, C::Query), C::Filter>>,
+) {
+    let mut values = Vec::with_capacity(*previous_len);
+    for (entity, computed_visibility, query_item) in query.iter_mut() {
+        if computed_visibility.is_visible {
+            values.push((entity, (C::extract_component(query_item),)));
+        }
+    }
+    *previous_len = values.len();
+    commands.insert_or_spawn_batch(values);
+}