Add dynamic vertex buffer layouts to resolve sorting issues

- Vertex buffer layouts are set on `Mesh` for shader specialization.
  `Mesh` knows how to specialize its own vertex buffer because it owns
  the vertex data.
- The `RenderAsset` implementation for `Mesh` caches the vertex buffer
  layout and adds a cache key to the `GpuMesh` render asset. The
  `MeshPipeline` can lookup the vertex buffer layout by composing
  `VertexLayoutKey` within `MeshPipelineKey`.
- `SpritePipeline` and `UiPipeline` populates the vertex layout cache
  for themselves..
- `SpecializedPipeline::specialize` now takes a reference to the
  `RenderPipelineCache`, which is the best way I could find to get a
  reference to the cache for vertex layout lookups.
- Discussion: https://discord.com/channels/691052431525675048/743663924229963868/908484759833960489

Author: parasyte

crates/bevy_gltf/src/loader.rs

@@ -14,6 +14,7 @@ use bevy_render::{
     color::Color,
     mesh::{Indices, Mesh, VertexAttributeValues},
     primitives::{Aabb, Frustum},
+    render_resource::VertexFormat,
     texture::{Image, ImageType, TextureError},
     view::VisibleEntities,
 };
@@ -133,6 +134,8 @@ async fn load_gltf<'a, 'b>(
                 .read_tangents()
                 .map(|v| VertexAttributeValues::Float32x4(v.collect()))
             {
+                mesh.vertex_layout_mut()
+                    .push(Mesh::ATTRIBUTE_TANGENT, VertexFormat::Float32x4);
                 mesh.set_attribute(Mesh::ATTRIBUTE_TANGENT, vertex_attribute);
             }
 

crates/bevy_pbr/src/render/light.rs

@@ -14,7 +14,7 @@ use bevy_math::{const_vec3, Mat4, UVec3, UVec4, Vec2, Vec3, Vec4, Vec4Swizzles};
 use bevy_render::{
     camera::{Camera, CameraProjection},
     color::Color,
-    mesh::Mesh,
+    mesh::{Mesh, VertexLayoutKey},
     render_asset::RenderAssets,
     render_graph::{Node, NodeRunError, RenderGraphContext, SlotInfo, SlotType},
     render_phase::{
@@ -207,84 +207,21 @@ impl FromWorld for ShadowPipeline {
     }
 }
 
-bitflags::bitflags! {
-    #[repr(transparent)]
-    pub struct ShadowPipelineKey: u32 {
-        const NONE               = 0;
-        const VERTEX_TANGENTS    = (1 << 0);
-    }
-}
+#[derive(Clone, Copy, Hash, PartialEq, Eq)]
+pub struct ShadowPipelineKey(pub VertexLayoutKey);
 
 impl SpecializedPipeline for ShadowPipeline {
     type Key = ShadowPipelineKey;
 
-    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
-        let (vertex_array_stride, vertex_attributes) =
-            if key.contains(ShadowPipelineKey::VERTEX_TANGENTS) {
-                (
-                    48,
-                    vec![
-                        // Position (GOTCHA! Vertex_Position isn't first in the buffer due to how Mesh sorts attributes (alphabetically))
-                        VertexAttribute {
-                            format: VertexFormat::Float32x3,
-                            offset: 12,
-                            shader_location: 0,
-                        },
-                        // Normal
-                        VertexAttribute {
-                            format: VertexFormat::Float32x3,
-                            offset: 0,
-                            shader_location: 1,
-                        },
-                        // Uv (GOTCHA! uv is no longer third in the buffer due to how Mesh sorts attributes (alphabetically))
-                        VertexAttribute {
-                            format: VertexFormat::Float32x2,
-                            offset: 40,
-                            shader_location: 2,
-                        },
-                        // Tangent
-                        VertexAttribute {
-                            format: VertexFormat::Float32x4,
-                            offset: 24,
-                            shader_location: 3,
-                        },
-                    ],
-                )
-            } else {
-                (
-                    32,
-                    vec![
-                        // Position (GOTCHA! Vertex_Position isn't first in the buffer due to how Mesh sorts attributes (alphabetically))
-                        VertexAttribute {
-                            format: VertexFormat::Float32x3,
-                            offset: 12,
-                            shader_location: 0,
-                        },
-                        // Normal
-                        VertexAttribute {
-                            format: VertexFormat::Float32x3,
-                            offset: 0,
-                            shader_location: 1,
-                        },
-                        // Uv
-                        VertexAttribute {
-                            format: VertexFormat::Float32x2,
-                            offset: 24,
-                            shader_location: 2,
-                        },
-                    ],
-                )
-            };
+    fn specialize(&self, cache: &RenderPipelineCache, key: Self::Key) -> RenderPipelineDescriptor {
+        let vertex_layout = cache.vertex_layout_cache.get(&key.0).unwrap();
+
         RenderPipelineDescriptor {
             vertex: VertexState {
                 shader: SHADOW_SHADER_HANDLE.typed::<Shader>(),
                 entry_point: "vertex".into(),
                 shader_defs: vec![],
-                buffers: vec![VertexBufferLayout {
-                    array_stride: vertex_array_stride,
-                    step_mode: VertexStepMode::Vertex,
-                    attributes: vertex_attributes,
-                }],
+                buffers: vec![vertex_layout.clone()],
             },
             fragment: None,
             layout: Some(vec![self.view_layout.clone(), self.mesh_layout.clone()]),
@@ -1063,13 +1000,9 @@ pub fn queue_shadows(
             // 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 mut key = ShadowPipelineKey::empty();
                 if let Ok(mesh_handle) = casting_meshes.get(entity) {
-                    if let Some(mesh) = render_meshes.get(mesh_handle) {
-                        if mesh.has_tangents {
-                            key |= ShadowPipelineKey::VERTEX_TANGENTS;
-                        }
-                    }
+                    let mesh = render_meshes.get(mesh_handle).unwrap();
+                    let key = ShadowPipelineKey(mesh.vertex_layout_key);
                     let pipeline_id =
                         pipelines.specialize(&mut pipeline_cache, &shadow_pipeline, key);
 

crates/bevy_pbr/src/render/mesh.rs

@@ -11,7 +11,7 @@ use bevy_ecs::{
 use bevy_math::Mat4;
 use bevy_reflect::TypeUuid;
 use bevy_render::{
-    mesh::Mesh,
+    mesh::{Mesh, VertexLayoutKey},
     render_asset::RenderAssets,
     render_component::{ComponentUniforms, DynamicUniformIndex, UniformComponentPlugin},
     render_phase::{EntityRenderCommand, RenderCommandResult, TrackedRenderPass},
@@ -361,25 +361,31 @@ impl MeshPipeline {
     }
 }
 
+#[derive(Clone, Copy, Hash, PartialEq, Eq)]
+pub struct MeshPipelineKey {
+    pub vertex_layout_key: VertexLayoutKey,
+    pub flags: MeshPipelineFlags,
+}
+
 bitflags::bitflags! {
     #[repr(transparent)]
     // NOTE: Apparently quadro drivers support up to 64x MSAA.
     /// MSAA uses the highest 6 bits for the MSAA sample count - 1 to support up to 64x MSAA.
-    pub struct MeshPipelineKey: u32 {
+    pub struct MeshPipelineFlags: u32 {
         const NONE                        = 0;
         const VERTEX_TANGENTS             = (1 << 0);
         const TRANSPARENT_MAIN_PASS       = (1 << 1);
-        const MSAA_RESERVED_BITS          = MeshPipelineKey::MSAA_MASK_BITS << MeshPipelineKey::MSAA_SHIFT_BITS;
+        const MSAA_RESERVED_BITS          = MeshPipelineFlags::MSAA_MASK_BITS << MeshPipelineFlags::MSAA_SHIFT_BITS;
     }
 }
 
-impl MeshPipelineKey {
+impl MeshPipelineFlags {
     const MSAA_MASK_BITS: u32 = 0b111111;
     const MSAA_SHIFT_BITS: u32 = 32 - 6;
 
     pub fn from_msaa_samples(msaa_samples: u32) -> Self {
         let msaa_bits = ((msaa_samples - 1) & Self::MSAA_MASK_BITS) << Self::MSAA_SHIFT_BITS;
-        MeshPipelineKey::from_bits(msaa_bits).unwrap()
+        MeshPipelineFlags::from_bits(msaa_bits).unwrap()
     }
 
     pub fn msaa_samples(&self) -> u32 {
@@ -390,70 +396,19 @@ impl MeshPipelineKey {
 impl SpecializedPipeline for MeshPipeline {
     type Key = MeshPipelineKey;
 
-    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
-        let (vertex_array_stride, vertex_attributes) =
-            if key.contains(MeshPipelineKey::VERTEX_TANGENTS) {
-                (
-                    48,
-                    vec![
-                        // Position (GOTCHA! Vertex_Position isn't first in the buffer due to how Mesh sorts attributes (alphabetically))
-                        VertexAttribute {
-                            format: VertexFormat::Float32x3,
-                            offset: 12,
-                            shader_location: 0,
-                        },
-                        // Normal
-                        VertexAttribute {
-                            format: VertexFormat::Float32x3,
-                            offset: 0,
-                            shader_location: 1,
-                        },
-                        // Uv (GOTCHA! uv is no longer third in the buffer due to how Mesh sorts attributes (alphabetically))
-                        VertexAttribute {
-                            format: VertexFormat::Float32x2,
-                            offset: 40,
-                            shader_location: 2,
-                        },
-                        // Tangent
-                        VertexAttribute {
-                            format: VertexFormat::Float32x4,
-                            offset: 24,
-                            shader_location: 3,
-                        },
-                    ],
-                )
-            } else {
-                (
-                    32,
-                    vec![
-                        // Position (GOTCHA! Vertex_Position isn't first in the buffer due to how Mesh sorts attributes (alphabetically))
-                        VertexAttribute {
-                            format: VertexFormat::Float32x3,
-                            offset: 12,
-                            shader_location: 0,
-                        },
-                        // Normal
-                        VertexAttribute {
-                            format: VertexFormat::Float32x3,
-                            offset: 0,
-                            shader_location: 1,
-                        },
-                        // Uv
-                        VertexAttribute {
-                            format: VertexFormat::Float32x2,
-                            offset: 24,
-                            shader_location: 2,
-                        },
-                    ],
-                )
-            };
+    fn specialize(&self, cache: &RenderPipelineCache, key: Self::Key) -> RenderPipelineDescriptor {
+        let vertex_layout = cache
+            .vertex_layout_cache
+            .get(&key.vertex_layout_key)
+            .unwrap();
+
         let mut shader_defs = Vec::new();
-        if key.contains(MeshPipelineKey::VERTEX_TANGENTS) {
+        if key.flags.contains(MeshPipelineFlags::VERTEX_TANGENTS) {
             shader_defs.push(String::from("VERTEX_TANGENTS"));
         }
 
         let (label, blend, depth_write_enabled);
-        if key.contains(MeshPipelineKey::TRANSPARENT_MAIN_PASS) {
+        if key.flags.contains(MeshPipelineFlags::TRANSPARENT_MAIN_PASS) {
             label = "transparent_mesh_pipeline".into();
             blend = Some(BlendState::ALPHA_BLENDING);
             // For the transparent pass, fragments that are closer will be alpha blended
@@ -473,11 +428,7 @@ impl SpecializedPipeline for MeshPipeline {
                 shader: MESH_SHADER_HANDLE.typed::<Shader>(),
                 entry_point: "vertex".into(),
                 shader_defs: shader_defs.clone(),
-                buffers: vec![VertexBufferLayout {
-                    array_stride: vertex_array_stride,
-                    step_mode: VertexStepMode::Vertex,
-                    attributes: vertex_attributes,
-                }],
+                buffers: vec![vertex_layout.clone()],
             },
             fragment: Some(FragmentState {
                 shader: MESH_SHADER_HANDLE.typed::<Shader>(),
@@ -516,7 +467,7 @@ impl SpecializedPipeline for MeshPipeline {
                 },
             }),
             multisample: MultisampleState {
-                count: key.msaa_samples(),
+                count: key.flags.msaa_samples(),
                 mask: !0,
                 alpha_to_coverage_enabled: false,
             },
@@ -710,11 +661,11 @@ impl EntityRenderCommand for DrawMesh {
 
 #[cfg(test)]
 mod tests {
-    use super::MeshPipelineKey;
+    use super::MeshPipelineFlags;
     #[test]
     fn mesh_key_msaa_samples() {
         for i in 1..=64 {
-            assert_eq!(MeshPipelineKey::from_msaa_samples(i).msaa_samples(), i);
+            assert_eq!(MeshPipelineFlags::from_msaa_samples(i).msaa_samples(), i);
         }
     }
 }

crates/bevy_pbr/src/render/mod.rs

@@ -191,8 +191,8 @@ pub struct PbrPipelineKey {
 impl SpecializedPipeline for PbrPipeline {
     type Key = PbrPipelineKey;
 
-    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
-        let mut descriptor = self.mesh_pipeline.specialize(key.mesh_key);
+    fn specialize(&self, cache: &RenderPipelineCache, key: Self::Key) -> RenderPipelineDescriptor {
+        let mut descriptor = self.mesh_pipeline.specialize(cache, key.mesh_key);
         descriptor.fragment.as_mut().unwrap().shader = PBR_SHADER_HANDLE.typed::<Shader>();
         descriptor.layout = Some(vec![
             self.mesh_pipeline.view_layout.clone(),
@@ -254,25 +254,24 @@ pub fn queue_meshes(
 
         let inverse_view_matrix = view.transform.compute_matrix().inverse();
         let inverse_view_row_2 = inverse_view_matrix.row(2);
-        let mesh_key = MeshPipelineKey::from_msaa_samples(msaa.samples);
+        let flags = MeshPipelineFlags::from_msaa_samples(msaa.samples);
 
         for visible_entity in &visible_entities.entities {
             if let Ok((material_handle, mesh_handle, mesh_uniform)) =
                 standard_material_meshes.get(*visible_entity)
             {
                 if let Some(material) = render_materials.get(material_handle) {
+                    let mesh = render_meshes.get(mesh_handle).unwrap();
                     let mut pbr_key = PbrPipelineKey {
-                        mesh_key,
+                        mesh_key: MeshPipelineKey {
+                            vertex_layout_key: mesh.vertex_layout_key,
+                            flags,
+                        },
                         normal_map: material.has_normal_map,
                     };
-                    if let Some(mesh) = render_meshes.get(mesh_handle) {
-                        if mesh.has_tangents {
-                            pbr_key.mesh_key |= MeshPipelineKey::VERTEX_TANGENTS;
-                        }
-                    }
 
                     if let AlphaMode::Blend = material.alpha_mode {
-                        pbr_key.mesh_key |= MeshPipelineKey::TRANSPARENT_MAIN_PASS
+                        pbr_key.mesh_key.flags |= MeshPipelineFlags::TRANSPARENT_MAIN_PASS
                     }
 
                     let pipeline_id =