From 31550e6a46a8f5c8f89ebf1a28cb0bf84af21593 Mon Sep 17 00:00:00 2001
From: Jay Oster <jay@kodewerx.org>
Date: Sat, 13 Nov 2021 09:49:39 -0800
Subject: [PATCH] Add a new vertex attribute to the pipelined custom shader
 example

---
 assets/shaders/custom_material.wgsl        | 36 ++++++++++++++-
 examples/shader/custom_shader_pipelined.rs | 54 ++++++++++++++++++++--
 2 files changed, 85 insertions(+), 5 deletions(-)

diff --git a/assets/shaders/custom_material.wgsl b/assets/shaders/custom_material.wgsl
index 9dfaa6f568f43c..75642b60df8a41 100644
--- a/assets/shaders/custom_material.wgsl
+++ b/assets/shaders/custom_material.wgsl
@@ -1,3 +1,31 @@
+#import bevy_pbr::mesh_view_bind_group
+#import bevy_pbr::mesh_struct
+
+struct Vertex {
+    [[location(0)]] position: vec3<f32>;
+    [[location(1)]] normal: vec3<f32>;
+    [[location(2)]] uv: vec2<f32>;
+    [[location(3)]] color: vec4<f32>;
+};
+
+struct VertexOutput {
+    [[builtin(position)]] clip_position: vec4<f32>;
+    [[location(0)]] color: vec4<f32>;
+};
+
+[[group(2), binding(0)]]
+var<uniform> mesh: Mesh;
+
+[[stage(vertex)]]
+fn vertex(vertex: Vertex) -> VertexOutput {
+    let world_position = mesh.model * vec4<f32>(vertex.position, 1.0);
+
+    var out: VertexOutput;
+    out.clip_position = view.view_proj * world_position;
+    out.color = vertex.color;
+    return out;
+}
+
 [[block]]
 struct CustomMaterial {
     color: vec4<f32>;
@@ -5,7 +33,11 @@ struct CustomMaterial {
 [[group(1), binding(0)]]
 var<uniform> material: CustomMaterial;
 
+struct FragmentInput {
+    [[location(0)]] color: vec4<f32>;
+};
+
 [[stage(fragment)]]
-fn fragment() -> [[location(0)]] vec4<f32> {
-    return material.color;
+fn fragment(fragment: FragmentInput) -> [[location(0)]] vec4<f32> {
+    return (material.color + fragment.color) / 2.0;
 }
diff --git a/examples/shader/custom_shader_pipelined.rs b/examples/shader/custom_shader_pipelined.rs
index 622c9d69565aef..8fe3d7c2339f44 100644
--- a/examples/shader/custom_shader_pipelined.rs
+++ b/examples/shader/custom_shader_pipelined.rs
@@ -1,11 +1,12 @@
 use bevy::{
+    core::Time,
     core_pipeline::Transparent3d,
     diagnostic::{FrameTimeDiagnosticsPlugin, LogDiagnosticsPlugin},
     ecs::{
         prelude::*,
         system::{lifetimeless::*, SystemParamItem},
     },
-    math::{Vec3, Vec4},
+    math::{Quat, Vec3, Vec4},
     pbr2::{
         DrawMesh, MeshPipeline, MeshPipelineKey, MeshUniform, SetMeshBindGroup,
         SetMeshViewBindGroup,
@@ -38,6 +39,7 @@ fn main() {
         .add_plugin(LogDiagnosticsPlugin::default())
         .add_plugin(CustomMaterialPlugin)
         .add_startup_system(setup)
+        .add_system(rotate_mesh)
         .run();
 }
 
@@ -47,9 +49,13 @@ fn setup(
     mut meshes: ResMut<Assets<Mesh>>,
     mut materials: ResMut<Assets<CustomMaterial>>,
 ) {
-    // cube
+    // cube with custom vertex attributes
+    let mut mesh = Mesh::from(shape::Cube { size: 1.0 });
+    mesh.vertex_layout_mut()
+        .push(Mesh::ATTRIBUTE_COLOR, VertexFormat::Float32x4);
+    mesh.set_attribute(Mesh::ATTRIBUTE_COLOR, cube_vertex_colors());
     commands.spawn().insert_bundle((
-        meshes.add(Mesh::from(shape::Cube { size: 1.0 })),
+        meshes.add(mesh),
         Transform::from_xyz(0.0, 0.5, 0.0),
         GlobalTransform::default(),
         Visibility::default(),
@@ -66,6 +72,48 @@ fn setup(
     });
 }
 
+fn rotate_mesh(time: Res<Time>, mut query: Query<&mut Transform, With<Handle<Mesh>>>) {
+    let angle = std::f32::consts::TAU / 12.0;
+    let mut transform = query.single_mut();
+    transform.rotate(Quat::from_rotation_x(angle * time.delta_seconds()));
+    transform.rotate(Quat::from_rotation_y(angle * time.delta_seconds()));
+}
+
+fn cube_vertex_colors() -> Vec<[f32; 4]> {
+    vec![
+        // Front
+        [0.0, 0.5, 0.0, 1.0], // Green
+        [0.5, 0.0, 0.0, 1.0], // Red
+        [0.0, 0.5, 0.5, 1.0], // Cyan
+        [0.5, 0.0, 0.5, 1.0], // Magenta
+        // Back
+        [0.5, 0.5, 0.0, 1.0], // Yellow
+        [0.0, 0.0, 0.0, 1.0], // Black
+        [0.5, 0.5, 0.5, 1.0], // White
+        [0.0, 0.0, 0.5, 1.0], // Blue
+        // Right
+        [0.5, 0.5, 0.5, 1.0], // White
+        [0.0, 0.0, 0.0, 1.0], // Black
+        [0.0, 0.5, 0.5, 1.0], // Cyan
+        [0.5, 0.0, 0.0, 1.0], // Red
+        // Left
+        [0.0, 0.5, 0.0, 1.0], // Green
+        [0.5, 0.0, 0.5, 1.0], // Magenta
+        [0.5, 0.5, 0.0, 1.0], // Yellow
+        [0.0, 0.0, 0.5, 1.0], // Blue
+        // Top
+        [0.0, 0.0, 0.0, 1.0], // Black
+        [0.5, 0.5, 0.0, 1.0], // Yellow
+        [0.5, 0.0, 0.5, 1.0], // Magenta
+        [0.0, 0.5, 0.5, 1.0], // Cyan
+        // Bottom
+        [0.5, 0.0, 0.0, 1.0], // Red
+        [0.0, 0.5, 0.0, 1.0], // Green
+        [0.0, 0.0, 0.5, 1.0], // Blue
+        [0.5, 0.5, 0.5, 1.0], // White
+    ]
+}
+
 #[derive(Debug, Clone, TypeUuid)]
 #[uuid = "4ee9c363-1124-4113-890e-199d81b00281"]
 pub struct CustomMaterial {