From 950abd79486621dc7c3677f799e40a9e5011a9cf Mon Sep 17 00:00:00 2001
From: Robert Swain <robert.swain@gmail.com>
Date: Mon, 4 Sep 2023 17:00:29 +0200
Subject: [PATCH] HACK: Split batches for skinning or morph targets

---
 crates/bevy_pbr/src/material.rs            | 21 +++++++++++++++------
 crates/bevy_pbr/src/render/mesh.rs         |  9 +++++++--
 crates/bevy_pbr/src/render/mesh_types.wgsl |  4 +++-
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/crates/bevy_pbr/src/material.rs b/crates/bevy_pbr/src/material.rs
index 4347c05629b0c..76eb8a1960ca2 100644
--- a/crates/bevy_pbr/src/material.rs
+++ b/crates/bevy_pbr/src/material.rs
@@ -1,7 +1,7 @@
 use crate::{
-    render, AlphaMode, DrawMesh, DrawPrepass, EnvironmentMapLight, MeshPipeline, MeshPipelineKey,
-    MeshTransforms, PrepassPipelinePlugin, PrepassPlugin, ScreenSpaceAmbientOcclusionSettings,
-    SetMeshBindGroup, SetMeshViewBindGroup, Shadow,
+    is_skinned, render, AlphaMode, DrawMesh, DrawPrepass, EnvironmentMapLight, MeshFlags,
+    MeshPipeline, MeshPipelineKey, MeshTransforms, PrepassPipelinePlugin, PrepassPlugin,
+    ScreenSpaceAmbientOcclusionSettings, SetMeshBindGroup, SetMeshViewBindGroup, Shadow,
 };
 use bevy_app::{App, Plugin};
 use bevy_asset::{AddAsset, AssetEvent, AssetServer, Assets, Handle};
@@ -409,7 +409,7 @@ pub fn queue_material_meshes<M: Material>(
         &Handle<M>,
         &mut MaterialBindGroupId,
         &Handle<Mesh>,
-        &MeshTransforms,
+        &mut MeshTransforms,
     )>,
     images: Res<RenderAssets<Image>>,
     mut views: Query<(
@@ -494,8 +494,12 @@ pub fn queue_material_meshes<M: Material>(
 
         let rangefinder = view.rangefinder3d();
         for visible_entity in &visible_entities.entities {
-            if let Ok((material_handle, mut material_bind_group_id, mesh_handle, mesh_transforms)) =
-                material_meshes.get_mut(*visible_entity)
+            if let Ok((
+                material_handle,
+                mut material_bind_group_id,
+                mesh_handle,
+                mut mesh_transforms,
+            )) = material_meshes.get_mut(*visible_entity)
             {
                 if let (Some(mesh), Some(material)) = (
                     render_meshes.get(mesh_handle),
@@ -504,8 +508,13 @@ pub fn queue_material_meshes<M: Material>(
                     let mut mesh_key =
                         MeshPipelineKey::from_primitive_topology(mesh.primitive_topology)
                             | view_key;
+
+                    if is_skinned(&mesh.layout) {
+                        mesh_transforms.flags |= MeshFlags::SKINNED.bits();
+                    }
                     if mesh.morph_targets.is_some() {
                         mesh_key |= MeshPipelineKey::MORPH_TARGETS;
+                        mesh_transforms.flags |= MeshFlags::MORPH_TARGETS.bits();
                     }
                     match material.properties.alpha_mode {
                         AlphaMode::Blend => {
diff --git a/crates/bevy_pbr/src/render/mesh.rs b/crates/bevy_pbr/src/render/mesh.rs
index e10a4a1dab164..de1049850c2c3 100644
--- a/crates/bevy_pbr/src/render/mesh.rs
+++ b/crates/bevy_pbr/src/render/mesh.rs
@@ -247,8 +247,10 @@ impl From<&MeshTransforms> for MeshUniform {
 // NOTE: These must match the bit flags in bevy_pbr/src/render/mesh_types.wgsl!
 bitflags::bitflags! {
     #[repr(transparent)]
-    struct MeshFlags: u32 {
+    pub struct MeshFlags: u32 {
         const SHADOW_RECEIVER            = (1 << 0);
+        const SKINNED                    = (1 << 1);
+        const MORPH_TARGETS              = (1 << 2);
         // Indicates the sign of the determinant of the 3x3 model matrix. If the sign is positive,
         // then the flag should be set, else it should not be set.
         const SIGN_DETERMINANT_MODEL_3X3 = (1 << 31);
@@ -413,6 +415,7 @@ struct BatchMeta<'mat, 'mesh> {
     /// dynamic offsets.
     material_binding_meta: Option<&'mat MaterialBindGroupId>,
     mesh_handle: Option<&'mesh Handle<Mesh>>,
+    mesh_flags: u32,
     dynamic_offset: Option<NonMaxU32>,
 }
 
@@ -422,6 +425,7 @@ impl<'mat, 'mesh> BatchMeta<'mat, 'mesh> {
         self.pipeline_id == other.pipeline_id
             && self.draw_function_id == other.draw_function_id
             && self.mesh_handle == other.mesh_handle
+            && (self.mesh_flags & (MeshFlags::SKINNED | MeshFlags::MORPH_TARGETS).bits()) == 0
             && self.dynamic_offset == other.dynamic_offset
             && (!consider_material || self.material_binding_meta == other.material_binding_meta)
     }
@@ -470,6 +474,7 @@ fn process_phase<I: CachedRenderPipelinePhaseItem>(
             draw_function_id: Some(item.draw_function()),
             material_binding_meta,
             mesh_handle: Some(mesh_handle),
+            mesh_flags: mesh_transforms.flags,
             dynamic_offset: gpu_array_buffer_index.dynamic_offset,
         };
         if !batch_meta.matches(&batch.meta, consider_material) {
@@ -939,7 +944,7 @@ impl MeshPipelineKey {
     }
 }
 
-fn is_skinned(layout: &Hashed<InnerMeshVertexBufferLayout>) -> bool {
+pub fn is_skinned(layout: &Hashed<InnerMeshVertexBufferLayout>) -> bool {
     layout.contains(Mesh::ATTRIBUTE_JOINT_INDEX) && layout.contains(Mesh::ATTRIBUTE_JOINT_WEIGHT)
 }
 pub fn setup_morph_and_skinning_defs(
diff --git a/crates/bevy_pbr/src/render/mesh_types.wgsl b/crates/bevy_pbr/src/render/mesh_types.wgsl
index 7412de7a8a5f7..ba04c18e4a7d6 100644
--- a/crates/bevy_pbr/src/render/mesh_types.wgsl
+++ b/crates/bevy_pbr/src/render/mesh_types.wgsl
@@ -28,6 +28,8 @@ struct MorphWeights {
 };
 #endif
 
-const MESH_FLAGS_SHADOW_RECEIVER_BIT: u32 = 1u;
+const MESH_FLAGS_SHADOW_RECEIVER_BIT: u32            = 1u;
+const MESH_FLAGS_SKINNED_BIT: u32                    = 2u;
+const MESH_FLAGS_MORPH_TARGETS_BIT: u32              = 4u;
 // 2^31 - if the flag is set, the sign is positive, else it is negative
 const MESH_FLAGS_SIGN_DETERMINANT_MODEL_3X3_BIT: u32 = 2147483648u;