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

examples/shader/custom_shader_pipelined.rs

@@ -7,7 +7,7 @@ use bevy::{
     },
     math::{Vec3, Vec4},
     pbr2::{
-        DrawMesh, MeshPipeline, MeshPipelineKey, MeshUniform, SetMeshBindGroup,
+        DrawMesh, MeshPipeline, MeshPipelineFlags, MeshPipelineKey, MeshUniform, SetMeshBindGroup,
         SetMeshViewBindGroup,
     },
     prelude::{AddAsset, App, AssetServer, Assets, GlobalTransform, Handle, Plugin, Transform},
@@ -135,8 +135,8 @@ pub struct CustomPipeline {
 impl SpecializedPipeline for CustomPipeline {
     type Key = MeshPipelineKey;
 
-    fn specialize(&self, key: Self::Key) -> RenderPipelineDescriptor {
-        let mut descriptor = self.mesh_pipeline.specialize(key);
+    fn specialize(&self, cache: &RenderPipelineCache, key: Self::Key) -> RenderPipelineDescriptor {
+        let mut descriptor = self.mesh_pipeline.specialize(cache, key);
         descriptor.fragment.as_mut().unwrap().shader = self.shader.clone();
         descriptor.layout = Some(vec![
             self.mesh_pipeline.view_layout.clone(),
@@ -177,22 +177,27 @@ impl FromWorld for CustomPipeline {
 pub fn queue_custom(
     transparent_3d_draw_functions: Res<DrawFunctions<Transparent3d>>,
     materials: Res<RenderAssets<CustomMaterial>>,
+    render_meshes: Res<RenderAssets<Mesh>>,
     custom_pipeline: Res<CustomPipeline>,
     mut pipeline_cache: ResMut<RenderPipelineCache>,
     mut specialized_pipelines: ResMut<SpecializedPipelines<CustomPipeline>>,
     msaa: Res<Msaa>,
-    material_meshes: Query<(Entity, &Handle<CustomMaterial>, &MeshUniform), With<Handle<Mesh>>>,
+    material_meshes: Query<(Entity, &Handle<Mesh>, &Handle<CustomMaterial>, &MeshUniform)>,
     mut views: Query<(&ExtractedView, &mut RenderPhase<Transparent3d>)>,
 ) {
     let draw_custom = transparent_3d_draw_functions
         .read()
         .get_id::<DrawCustom>()
         .unwrap();
-    let key = MeshPipelineKey::from_msaa_samples(msaa.samples);
     for (view, mut transparent_phase) in views.iter_mut() {
         let view_matrix = view.transform.compute_matrix();
         let view_row_2 = view_matrix.row(2);
-        for (entity, material_handle, mesh_uniform) in material_meshes.iter() {
+        for (entity, mesh_handle, material_handle, mesh_uniform) in material_meshes.iter() {
+            let mesh = render_meshes.get(mesh_handle).unwrap();
+            let key = MeshPipelineKey {
+                vertex_layout_key: mesh.vertex_layout_key,
+                flags: MeshPipelineFlags::from_msaa_samples(msaa.samples),
+            };
             if materials.contains_key(material_handle) {
                 transparent_phase.add(Transparent3d {
                     entity,

examples/shader/shader_defs_pipelined.rs

@@ -4,13 +4,14 @@ use bevy::{
     ecs::prelude::*,
     math::Vec3,
     pbr2::{
-        DrawMesh, MeshPipeline, MeshPipelineKey, MeshUniform, SetMeshBindGroup,
+        DrawMesh, MeshPipeline, MeshPipelineFlags, MeshPipelineKey, MeshUniform, SetMeshBindGroup,
         SetMeshViewBindGroup,
     },
     prelude::{App, AssetServer, Assets, GlobalTransform, Handle, Plugin, Transform},
     render2::{
         camera::PerspectiveCameraBundle,
         mesh::{shape, Mesh},
+        render_asset::RenderAssets,
         render_component::{ExtractComponent, ExtractComponentPlugin},
         render_phase::{AddRenderCommand, DrawFunctions, RenderPhase, SetItemPipeline},
         render_resource::*,
@@ -86,7 +87,7 @@ fn setup(mut commands: Commands, mut meshes: ResMut<Assets<Mesh>>) {
 }
 
 struct IsRedPipeline {
-    mesh_pipline: MeshPipeline,
+    mesh_pipeline: MeshPipeline,
     shader: Handle<Shader>,
 }
 
@@ -96,7 +97,7 @@ impl FromWorld for IsRedPipeline {
         let mesh_pipeline = world.get_resource::<MeshPipeline>().unwrap();
         let shader = asset_server.load("shaders/shader_defs.wgsl");
         IsRedPipeline {
-            mesh_pipline: mesh_pipeline.clone(),
+            mesh_pipeline: mesh_pipeline.clone(),
             shader,
         }
     }
@@ -105,20 +106,24 @@ impl FromWorld for IsRedPipeline {
 impl SpecializedPipeline for IsRedPipeline {
     type Key = (IsRed, MeshPipelineKey);
 
-    fn specialize(&self, (is_red, pbr_pipeline_key): Self::Key) -> RenderPipelineDescriptor {
+    fn specialize(
+        &self,
+        cache: &RenderPipelineCache,
+        (is_red, pbr_pipeline_key): Self::Key,
+    ) -> RenderPipelineDescriptor {
         let mut shader_defs = Vec::new();
         if is_red.0 {
             shader_defs.push("IS_RED".to_string());
         }
-        let mut descriptor = self.mesh_pipline.specialize(pbr_pipeline_key);
+        let mut descriptor = self.mesh_pipeline.specialize(cache, pbr_pipeline_key);
         descriptor.vertex.shader = self.shader.clone();
         descriptor.vertex.shader_defs = shader_defs.clone();
         let fragment = descriptor.fragment.as_mut().unwrap();
         fragment.shader = self.shader.clone();
         fragment.shader_defs = shader_defs;
         descriptor.layout = Some(vec![
-            self.mesh_pipline.view_layout.clone(),
-            self.mesh_pipline.mesh_layout.clone(),
+            self.mesh_pipeline.view_layout.clone(),
+            self.mesh_pipeline.mesh_layout.clone(),
         ]);
         descriptor
     }
@@ -131,24 +136,30 @@ type DrawIsRed = (
     DrawMesh,
 );
 
+#[allow(clippy::too_many_arguments)]
 fn queue_custom(
     transparent_3d_draw_functions: Res<DrawFunctions<Transparent3d>>,
+    render_meshes: Res<RenderAssets<Mesh>>,
     custom_pipeline: Res<IsRedPipeline>,
     msaa: Res<Msaa>,
     mut pipelines: ResMut<SpecializedPipelines<IsRedPipeline>>,
     mut pipeline_cache: ResMut<RenderPipelineCache>,
-    material_meshes: Query<(Entity, &MeshUniform, &IsRed), With<Handle<Mesh>>>,
+    material_meshes: Query<(Entity, &Handle<Mesh>, &MeshUniform, &IsRed)>,
     mut views: Query<(&ExtractedView, &mut RenderPhase<Transparent3d>)>,
 ) {
     let draw_custom = transparent_3d_draw_functions
         .read()
         .get_id::<DrawIsRed>()
         .unwrap();
-    let key = MeshPipelineKey::from_msaa_samples(msaa.samples);
     for (view, mut transparent_phase) in views.iter_mut() {
         let view_matrix = view.transform.compute_matrix();
         let view_row_2 = view_matrix.row(2);
-        for (entity, mesh_uniform, is_red) in material_meshes.iter() {
+        for (entity, mesh_handle, mesh_uniform, is_red) in material_meshes.iter() {
+            let mesh = render_meshes.get(mesh_handle).unwrap();
+            let key = MeshPipelineKey {
+                vertex_layout_key: mesh.vertex_layout_key,
+                flags: MeshPipelineFlags::from_msaa_samples(msaa.samples),
+            };
             let pipeline =
                 pipelines.specialize(&mut pipeline_cache, &custom_pipeline, (*is_red, key));
             transparent_phase.add(Transparent3d {

pipelined/bevy_gltf2/src/loader.rs

@@ -14,6 +14,7 @@ use bevy_render2::{
     color::Color,
     mesh::{Indices, Mesh, VertexAttributeValues},
     primitives::Aabb,
+    render_resource::VertexFormat,
     texture::{Image, ImageType, TextureError},
 };
 use bevy_scene::Scene;
@@ -132,6 +133,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);
             }
 

pipelined/bevy_pbr2/src/render/light.rs

@@ -14,7 +14,7 @@ use bevy_math::{const_vec3, Mat4, UVec3, UVec4, Vec3, Vec4, Vec4Swizzles};
 use bevy_render2::{
     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()]),
@@ -1040,13 +977,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);