From fe419295a3c523b36dcddfe4a414f3b9a30369c6 Mon Sep 17 00:00:00 2001
From: Robert Swain <robert.swain@gmail.com>
Date: Sun, 30 Jul 2023 21:55:01 +0200
Subject: [PATCH] Fix shader_material_glsl example after #9254

---
 assets/shaders/custom_material.vert | 14 ++++++++++++--
 crates/bevy_pbr/src/render/mesh.rs  | 25 +++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/assets/shaders/custom_material.vert b/assets/shaders/custom_material.vert
index 59af5305ba7b2a..f8d6b128a9ae59 100644
--- a/assets/shaders/custom_material.vert
+++ b/assets/shaders/custom_material.vert
@@ -16,13 +16,23 @@ layout(set = 0, binding = 0) uniform CameraViewProj {
     float height;
 };
 
-layout(set = 2, binding = 0) uniform Mesh {
+struct Mesh {
     mat4 Model;
     mat4 InverseTransposeModel;
     uint flags;
 };
 
+#ifdef PER_OBJECT_BUFFER_BATCH_SIZE
+layout(set = 2, binding = 0) uniform Mesh Meshes[#{PER_OBJECT_BUFFER_BATCH_SIZE}];
+#else
+layout(set = 2, binding = 0) readonly buffer _Meshes {
+    Mesh Meshes[];
+};
+#endif // PER_OBJECT_BUFFER_BATCH_SIZE
+
 void main() {
     v_Uv = Vertex_Uv;
-    gl_Position = ViewProj * Model * vec4(Vertex_Position, 1.0);
+    gl_Position = ViewProj
+        * Meshes[gl_BaseInstance + gl_InstanceIndex].Model
+        * vec4(Vertex_Position, 1.0);
 }
diff --git a/crates/bevy_pbr/src/render/mesh.rs b/crates/bevy_pbr/src/render/mesh.rs
index 97a45e4151c8ec..2dee3a39e6e963 100644
--- a/crates/bevy_pbr/src/render/mesh.rs
+++ b/crates/bevy_pbr/src/render/mesh.rs
@@ -325,6 +325,20 @@ pub struct MeshPipeline {
     pub dummy_white_gpu_image: GpuImage,
     pub clustered_forward_buffer_binding_type: BufferBindingType,
     pub mesh_layouts: MeshLayouts,
+    /// `MeshUniform`s are stored in arrays in buffers. If storage buffers are available, they
+    /// are used and this will be `None`, otherwise uniform buffers will be used with batches
+    /// of this many `MeshUniform`s, stored at dynamic offsets within the uniform buffer.
+    /// Use code like this in custom shaders:
+    /// ```
+    /// #ifdef PER_OBJECT_BUFFER_BATCH_SIZE
+    /// @group(2) @binding(0)
+    /// var<uniform> mesh: array<Mesh, #{PER_OBJECT_BUFFER_BATCH_SIZE}u>;
+    /// #else
+    /// @group(2) @binding(0)
+    /// var<storage> mesh: array<Mesh>;
+    /// #endif // PER_OBJECT_BUFFER_BATCH_SIZE
+    /// ```
+    pub per_object_buffer_batch_size: Option<u32>,
 }
 
 impl FromWorld for MeshPipeline {
@@ -564,6 +578,7 @@ impl FromWorld for MeshPipeline {
             clustered_forward_buffer_binding_type,
             dummy_white_gpu_image,
             mesh_layouts: MeshLayouts::new(&render_device),
+            per_object_buffer_batch_size: GpuArrayBuffer::<MeshUniform>::batch_size(&render_device),
         }
     }
 }
@@ -866,6 +881,16 @@ impl SpecializedMeshPipeline for MeshPipeline {
             TextureFormat::bevy_default()
         };
 
+        // This is defined here so that custom shaders that use something other than
+        // the mesh binding from bevy_pbr::mesh_bindings can easily make use of this
+        // in their own shaders.
+        if let Some(per_object_buffer_batch_size) = self.per_object_buffer_batch_size {
+            shader_defs.push(ShaderDefVal::UInt(
+                "PER_OBJECT_BUFFER_BATCH_SIZE".into(),
+                per_object_buffer_batch_size,
+            ));
+        }
+
         Ok(RenderPipelineDescriptor {
             vertex: VertexState {
                 shader: MESH_SHADER_HANDLE.typed::<Shader>(),