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

@@ -18,7 +18,7 @@ use bevy_render::{
     mesh::{Indices, Mesh, VertexAttributeValues},
     primitives::{Aabb, Frustum},
     render_resource::{
-        AddressMode, FilterMode, PrimitiveTopology, SamplerDescriptor, TextureFormat,
+        AddressMode, FilterMode, PrimitiveTopology, SamplerDescriptor, TextureFormat, VertexFormat,
     },
     texture::{Image, ImageType, TextureError},
     view::VisibleEntities,
@@ -138,6 +138,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/material.rs

@@ -1,6 +1,6 @@
 use crate::{
-    AlphaMode, DrawMesh, MeshPipeline, MeshPipelineKey, MeshUniform, SetMeshBindGroup,
-    SetMeshViewBindGroup,
+    AlphaMode, DrawMesh, MeshPipeline, MeshPipelineFlags, MeshPipelineKey, MeshUniform,
+    SetMeshBindGroup, SetMeshViewBindGroup,
 };
 use bevy_app::{App, Plugin};
 use bevy_asset::{AddAsset, Asset, AssetServer, Handle};
@@ -205,8 +205,8 @@ pub struct MaterialPipeline<M: SpecializedMaterial> {
 impl<M: SpecializedMaterial> SpecializedPipeline for MaterialPipeline<M> {
     type Key = (MeshPipelineKey, M::Key);
 
-    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
-        let mut descriptor = self.mesh_pipeline.specialize(key.0);
+    fn specialize(&self, cache: &RenderPipelineCache, key: Self::Key) -> RenderPipelineDescriptor {
+        let mut descriptor = self.mesh_pipeline.specialize(cache, key.0);
         if let Some(vertex_shader) = &self.vertex_shader {
             descriptor.vertex.shader = vertex_shader.clone();
         }
@@ -307,27 +307,33 @@ pub fn queue_material_meshes<M: SpecializedMaterial>(
 
         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 mesh_flags = MeshPipelineFlags::from_msaa_samples(msaa.samples);
 
         for visible_entity in &visible_entities.entities {
             if let Ok((material_handle, mesh_handle, mesh_uniform)) =
                 material_meshes.get(*visible_entity)
             {
-                if let Some(material) = render_materials.get(material_handle) {
-                    let mut mesh_key = mesh_key;
-                    if let Some(mesh) = render_meshes.get(mesh_handle) {
-                        if mesh.has_tangents {
-                            mesh_key |= MeshPipelineKey::VERTEX_TANGENTS;
-                        }
-                        mesh_key |=
-                            MeshPipelineKey::from_primitive_topology(mesh.primitive_topology);
+                if let (Some(material), Some(mesh)) = (
+                    render_materials.get(material_handle),
+                    render_meshes.get(mesh_handle),
+                ) {
+                    let mut mesh_flags = mesh_flags;
+                    if mesh.has_tangents {
+                        mesh_flags |= MeshPipelineFlags::VERTEX_TANGENTS;
                     }
+                    mesh_flags |=
+                        MeshPipelineFlags::from_primitive_topology(mesh.primitive_topology);
                     let alpha_mode = M::alpha_mode(material);
                     if let AlphaMode::Blend = alpha_mode {
-                        mesh_key |= MeshPipelineKey::TRANSPARENT_MAIN_PASS
+                        mesh_flags |= MeshPipelineFlags::TRANSPARENT_MAIN_PASS;
                     }
 
+                    let mesh_key = MeshPipelineKey {
+                        vertex_layout_key: mesh.vertex_layout_key,
+                        flags: mesh_flags,
+                    };
                     let specialized_key = M::key(material);
+
                     let pipeline_id = pipelines.specialize(
                         &mut pipeline_cache,
                         &material_pipeline,

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::{
@@ -209,84 +209,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()]),
@@ -1087,13 +1024,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::{GpuBufferInfo, Mesh},
+    mesh::{GpuBufferInfo, Mesh, VertexLayoutKey},
     render_asset::RenderAssets,
     render_component::{ComponentUniforms, DynamicUniformIndex, UniformComponentPlugin},
     render_phase::{EntityRenderCommand, RenderCommandResult, TrackedRenderPass},
@@ -356,28 +356,34 @@ 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 PRIMITIVE_TOPOLOGY_RESERVED_BITS = MeshPipelineKey::PRIMITIVE_TOPOLOGY_MASK_BITS << MeshPipelineKey::PRIMITIVE_TOPOLOGY_SHIFT_BITS;
+        const MSAA_RESERVED_BITS          = MeshPipelineFlags::MSAA_MASK_BITS << MeshPipelineFlags::MSAA_SHIFT_BITS;
+        const PRIMITIVE_TOPOLOGY_RESERVED_BITS = MeshPipelineFlags::PRIMITIVE_TOPOLOGY_MASK_BITS << MeshPipelineFlags::PRIMITIVE_TOPOLOGY_SHIFT_BITS;
     }
 }
 
-impl MeshPipelineKey {
+impl MeshPipelineFlags {
     const MSAA_MASK_BITS: u32 = 0b111111;
     const MSAA_SHIFT_BITS: u32 = 32 - 6;
     const PRIMITIVE_TOPOLOGY_MASK_BITS: u32 = 0b111;
     const PRIMITIVE_TOPOLOGY_SHIFT_BITS: u32 = Self::MSAA_SHIFT_BITS - 3;
 
     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 {
@@ -388,7 +394,7 @@ impl MeshPipelineKey {
         let primitive_topology_bits = ((primitive_topology as u32)
             & Self::PRIMITIVE_TOPOLOGY_MASK_BITS)
             << Self::PRIMITIVE_TOPOLOGY_SHIFT_BITS;
-        MeshPipelineKey::from_bits(primitive_topology_bits).unwrap()
+        MeshPipelineFlags::from_bits(primitive_topology_bits).unwrap()
     }
 
     pub fn primitive_topology(&self) -> PrimitiveTopology {
@@ -408,70 +414,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
@@ -494,11 +449,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>(),
@@ -517,7 +468,7 @@ impl SpecializedPipeline for MeshPipeline {
                 unclipped_depth: false,
                 polygon_mode: PolygonMode::Fill,
                 conservative: false,
-                topology: key.primitive_topology(),
+                topology: key.flags.primitive_topology(),
                 strip_index_format: None,
             },
             depth_stencil: Some(DepthStencilState {
@@ -537,7 +488,7 @@ impl SpecializedPipeline for MeshPipeline {
                 },
             }),
             multisample: MultisampleState {
-                count: key.msaa_samples(),
+                count: key.flags.msaa_samples(),
                 mask: !0,
                 alpha_to_coverage_enabled: false,
             },
@@ -737,11 +688,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);
         }
     }
 }