From 48fc396f04b04f7ef693e19ce3939e01ade2ef92 Mon Sep 17 00:00:00 2001
From: charlotte <charlotte.c.mcelwain@gmail.com>
Date: Tue, 18 Mar 2025 23:22:39 -0700
Subject: [PATCH] Fix unecessary specialization checks for apps with many
 materials (#18410)

# Objective

For materials that aren't being used or a visible entity doesn't have an
instance of, we were unnecessarily constantly checking whether they
needed specialization, saying yes (because the material had never been
specialized for that entity), and failing to look up the material
instance.

## Solution

If an entity doesn't have an instance of the material, it can't possibly
need specialization, so exit early before spending time doing the check.

Fixes #18388.
---
 crates/bevy_pbr/src/material.rs           |  6 +++---
 crates/bevy_pbr/src/prepass/mod.rs        |  7 +++----
 crates/bevy_pbr/src/render/light.rs       | 13 ++++++-------
 crates/bevy_sprite/src/mesh2d/material.rs |  7 +++----
 4 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/crates/bevy_pbr/src/material.rs b/crates/bevy_pbr/src/material.rs
index 1ac3930a0a..34ffb930db 100644
--- a/crates/bevy_pbr/src/material.rs
+++ b/crates/bevy_pbr/src/material.rs
@@ -863,6 +863,9 @@ pub fn specialize_material_meshes<M: Material>(
             .or_default();
 
         for (_, visible_entity) in visible_entities.iter::<Mesh3d>() {
+            let Some(material_asset_id) = render_material_instances.get(visible_entity) else {
+                continue;
+            };
             let entity_tick = entity_specialization_ticks.get(visible_entity).unwrap();
             let last_specialized_tick = view_specialized_material_pipeline_cache
                 .get(visible_entity)
@@ -874,9 +877,6 @@ pub fn specialize_material_meshes<M: Material>(
             if !needs_specialization {
                 continue;
             }
-            let Some(material_asset_id) = render_material_instances.get(visible_entity) else {
-                continue;
-            };
             let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(*visible_entity)
             else {
                 continue;
diff --git a/crates/bevy_pbr/src/prepass/mod.rs b/crates/bevy_pbr/src/prepass/mod.rs
index ea60f42667..724e96b429 100644
--- a/crates/bevy_pbr/src/prepass/mod.rs
+++ b/crates/bevy_pbr/src/prepass/mod.rs
@@ -962,6 +962,9 @@ pub fn specialize_prepass_material_meshes<M>(
             .or_default();
 
         for (_, visible_entity) in visible_entities.iter::<Mesh3d>() {
+            let Some(material_asset_id) = render_material_instances.get(visible_entity) else {
+                continue;
+            };
             let entity_tick = entity_specialization_ticks.get(visible_entity).unwrap();
             let last_specialized_tick = view_specialized_material_pipeline_cache
                 .get(visible_entity)
@@ -973,10 +976,6 @@ pub fn specialize_prepass_material_meshes<M>(
             if !needs_specialization {
                 continue;
             }
-
-            let Some(material_asset_id) = render_material_instances.get(visible_entity) else {
-                continue;
-            };
             let Some(mesh_instance) = render_mesh_instances.render_mesh_queue_data(*visible_entity)
             else {
                 continue;
diff --git a/crates/bevy_pbr/src/render/light.rs b/crates/bevy_pbr/src/render/light.rs
index 6f3216d896..a770a83bca 100644
--- a/crates/bevy_pbr/src/render/light.rs
+++ b/crates/bevy_pbr/src/render/light.rs
@@ -1809,6 +1809,9 @@ pub fn specialize_shadows<M: Material>(
                 .or_default();
 
             for (_, visible_entity) in visible_entities.iter().copied() {
+                let Some(material_asset_id) = render_material_instances.get(&visible_entity) else {
+                    continue;
+                };
                 let entity_tick = entity_specialization_ticks.get(&visible_entity).unwrap();
                 let last_specialized_tick = view_specialized_material_pipeline_cache
                     .get(&visible_entity)
@@ -1820,7 +1823,9 @@ pub fn specialize_shadows<M: Material>(
                 if !needs_specialization {
                     continue;
                 }
-
+                let Some(material) = render_materials.get(*material_asset_id) else {
+                    continue;
+                };
                 let Some(mesh_instance) =
                     render_mesh_instances.render_mesh_queue_data(visible_entity)
                 else {
@@ -1832,12 +1837,6 @@ pub fn specialize_shadows<M: Material>(
                 {
                     continue;
                 }
-                let Some(material_asset_id) = render_material_instances.get(&visible_entity) else {
-                    continue;
-                };
-                let Some(material) = render_materials.get(*material_asset_id) else {
-                    continue;
-                };
                 let Some(material_bind_group) =
                     material_bind_group_allocator.get(material.binding.group)
                 else {
diff --git a/crates/bevy_sprite/src/mesh2d/material.rs b/crates/bevy_sprite/src/mesh2d/material.rs
index 81ad7a9e3e..ce2803b261 100644
--- a/crates/bevy_sprite/src/mesh2d/material.rs
+++ b/crates/bevy_sprite/src/mesh2d/material.rs
@@ -698,6 +698,9 @@ pub fn specialize_material2d_meshes<M: Material2d>(
             .or_default();
 
         for (_, visible_entity) in visible_entities.iter::<Mesh2d>() {
+            let Some(material_asset_id) = render_material_instances.get(visible_entity) else {
+                continue;
+            };
             let entity_tick = entity_specialization_ticks.get(visible_entity).unwrap();
             let last_specialized_tick = view_specialized_material_pipeline_cache
                 .get(visible_entity)
@@ -709,10 +712,6 @@ pub fn specialize_material2d_meshes<M: Material2d>(
             if !needs_specialization {
                 continue;
             }
-
-            let Some(material_asset_id) = render_material_instances.get(visible_entity) else {
-                continue;
-            };
             let Some(mesh_instance) = render_mesh_instances.get_mut(visible_entity) else {
                 continue;
             };