From 479f43bbf34834ad2d4667de43351b6fa51f22d1 Mon Sep 17 00:00:00 2001
From: Robert Swain <robert.swain@gmail.com>
Date: Tue, 3 May 2022 18:28:04 +0000
Subject: [PATCH] Filter material handles on extraction (#4178)

# Objective

- While optimising many_cubes, I noticed that all material handles are extracted regardless of whether the entity to which the handle belongs is visible or not. As such >100k handles are extracted when only <20k are visible.

## Solution

- Only extract material handles of visible entities.
- This improves `many_cubes -- sphere` from ~42fps to ~48fps. It reduces not only the extraction time but also system commands time. `Handle<StandardMaterial>` extraction and its system commands went from 0.522ms + 3.710ms respectively, to 0.267ms + 0.227ms an 88% reduction for this system for this case. It's very view dependent but...
---
 crates/bevy_pbr/src/material.rs            |  2 +-
 crates/bevy_render/src/render_component.rs | 43 ++++++++++++++++++++--
 crates/bevy_sprite/src/mesh2d/material.rs  |  2 +-
 3 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/crates/bevy_pbr/src/material.rs b/crates/bevy_pbr/src/material.rs
index 2b1fd83cabe7d..b0d20dde6e25a 100644
--- a/crates/bevy_pbr/src/material.rs
+++ b/crates/bevy_pbr/src/material.rs
@@ -204,7 +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>>::default())
+            .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
diff --git a/crates/bevy_render/src/render_component.rs b/crates/bevy_render/src/render_component.rs
index 8d15f2476d6ae..7cd3158cb94d7 100644
--- 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);
+}
diff --git a/crates/bevy_sprite/src/mesh2d/material.rs b/crates/bevy_sprite/src/mesh2d/material.rs
index f088e761a054c..77cb3c4ebe92f 100644
--- a/crates/bevy_sprite/src/mesh2d/material.rs
+++ b/crates/bevy_sprite/src/mesh2d/material.rs
@@ -194,7 +194,7 @@ impl<M: SpecializedMaterial2d> Default for Material2dPlugin<M> {
 impl<M: SpecializedMaterial2d> Plugin for Material2dPlugin<M> {
     fn build(&self, app: &mut App) {
         app.add_asset::<M>()
-            .add_plugin(ExtractComponentPlugin::<Handle<M>>::default())
+            .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