From 5a80678502681a6212af12a0dbd7d916e209cdc0 Mon Sep 17 00:00:00 2001
From: robtfm <50659922+robtfm@users.noreply.github.com>
Date: Thu, 1 Feb 2024 09:46:28 +0000
Subject: [PATCH] fix shadow batching

---
 crates/bevy_pbr/src/render/light.rs | 6 ++++--
 crates/bevy_pbr/src/render/mesh.rs  | 8 +++++++-
 2 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/crates/bevy_pbr/src/render/light.rs b/crates/bevy_pbr/src/render/light.rs
index 8c2294c66c692..5abc500b109d8 100644
--- a/crates/bevy_pbr/src/render/light.rs
+++ b/crates/bevy_pbr/src/render/light.rs
@@ -1590,7 +1590,7 @@ pub fn queue_shadows<M: Material>(
     shadow_draw_functions: Res<DrawFunctions<Shadow>>,
     prepass_pipeline: Res<PrepassPipeline<M>>,
     render_meshes: Res<RenderAssets<Mesh>>,
-    render_mesh_instances: Res<RenderMeshInstances>,
+    mut render_mesh_instances: ResMut<RenderMeshInstances>,
     render_materials: Res<RenderMaterials<M>>,
     render_material_instances: Res<RenderMaterialInstances<M>>,
     mut pipelines: ResMut<SpecializedMeshPipelines<PrepassPipeline<M>>>,
@@ -1635,7 +1635,7 @@ pub fn queue_shadows<M: Material>(
             // NOTE: Lights with shadow mapping disabled will have no visible entities
             // so no meshes will be queued
             for entity in visible_entities.iter().copied() {
-                let Some(mesh_instance) = render_mesh_instances.get(&entity) else {
+                let Some(mesh_instance) = render_mesh_instances.get_mut(&entity) else {
                     continue;
                 };
                 if !mesh_instance.shadow_caster {
@@ -1685,6 +1685,8 @@ pub fn queue_shadows<M: Material>(
                     }
                 };
 
+                mesh_instance.material_bind_group_id = material.get_bind_group_id();
+
                 shadow_phase.add(Shadow {
                     draw_function: draw_shadow_mesh,
                     pipeline: pipeline_id,
diff --git a/crates/bevy_pbr/src/render/mesh.rs b/crates/bevy_pbr/src/render/mesh.rs
index 172a032b0401e..0722c1cf03df4 100644
--- a/crates/bevy_pbr/src/render/mesh.rs
+++ b/crates/bevy_pbr/src/render/mesh.rs
@@ -261,6 +261,12 @@ pub struct RenderMeshInstance {
     pub automatic_batching: bool,
 }
 
+impl RenderMeshInstance {
+    pub fn should_batch(&self) -> bool {
+        self.automatic_batching && self.material_bind_group_id.is_some()
+    }
+}
+
 #[derive(Default, Resource, Deref, DerefMut)]
 pub struct RenderMeshInstances(EntityHashMap<Entity, RenderMeshInstance>);
 
@@ -494,7 +500,7 @@ impl GetBatchData for MeshPipeline {
                 &mesh_instance.transforms,
                 maybe_lightmap.map(|lightmap| lightmap.uv_rect),
             ),
-            mesh_instance.automatic_batching.then_some((
+            mesh_instance.should_batch().then_some((
                 mesh_instance.material_bind_group_id,
                 mesh_instance.mesh_asset_id,
                 maybe_lightmap.map(|lightmap| lightmap.image),