Add reusable shader functions for transforming position/normal/tangent (#4901)

# Objective

- Add reusable shader functions for transforming positions / normals / tangents between local and world / clip space for 2D and 3D so that they are done in a simple and correct way
- The next step in #3969 so check there for more details.

## Solution

- Add `bevy_pbr::mesh_functions` and `bevy_sprite::mesh2d_functions` shader imports
  - These contain `mesh_` and `mesh2d_` versions of the following functions:
    - `mesh_position_local_to_world`
    - `mesh_position_world_to_clip`
    - `mesh_position_local_to_clip`
    - `mesh_normal_local_to_world`
    - `mesh_tangent_local_to_world`
- Use them everywhere where it is appropriate
  - Notably not in the sprite and UI shaders where `mesh2d_position_world_to_clip` could have been used, but including all the functions depends on the mesh binding so I chose to not use the function there
- NOTE: The `mesh_` and `mesh2d_` functions are currently identical. However, if I had defined only `bevy_pbr::mesh_functions` and used that in bevy_sprite, then bevy_sprite would have a runtime dependency on bevy_pbr, which seems undesirable. I also expect that when we have a proper 2D rendering API, these functions will diverge between 2D and 3D.

---

## Changelog

- Added: `bevy_pbr::mesh_functions` and `bevy_sprite::mesh2d_functions` shader imports containing `mesh_` and `mesh2d_` versions of the following functions:
  - `mesh_position_local_to_world`
  - `mesh_position_world_to_clip`
  - `mesh_position_local_to_clip`
  - `mesh_normal_local_to_world`
  - `mesh_tangent_local_to_world`

## Migration Guide

- The `skin_tangents` function from the `bevy_pbr::skinning` shader import has been replaced with the `mesh_tangent_local_to_world` function from the `bevy_pbr::mesh_functions` shader import

assets/shaders/animate_shader.wgsl

@@ -4,6 +4,9 @@
 [[group(1), binding(0)]]
 var<uniform> mesh: Mesh;
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
     [[location(1)]] normal: vec3<f32>;
@@ -17,10 +20,8 @@ struct VertexOutput {
 
 [[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.clip_position = mesh_position_local_to_clip(mesh.model, vec4<f32>(vertex.position, 1.0));
     out.uv = vertex.uv;
     return out;
 }

assets/shaders/custom_vertex_attribute.wgsl

@@ -1,17 +1,20 @@
 #import bevy_pbr::mesh_view_bindings
 #import bevy_pbr::mesh_bindings
 
-struct Vertex {
-    [[location(0)]] position: vec3<f32>;
-    [[location(1)]] blend_color: vec4<f32>;
-};
-
 struct CustomMaterial {
     color: vec4<f32>;
 };
 [[group(1), binding(0)]]
 var<uniform> material: CustomMaterial;
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
+struct Vertex {
+    [[location(0)]] position: vec3<f32>;
+    [[location(1)]] blend_color: vec4<f32>;
+};
+
 struct VertexOutput {
     [[builtin(position)]] clip_position: vec4<f32>;
     [[location(0)]] blend_color: vec4<f32>;
@@ -19,10 +22,8 @@ struct VertexOutput {
 
 [[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.clip_position = mesh_position_local_to_clip(mesh.model, vec4<f32>(vertex.position, 1.0));
     out.blend_color = vertex.blend_color;
     return out;
 }

assets/shaders/instancing.wgsl

@@ -4,6 +4,9 @@
 [[group(1), binding(0)]]
 var<uniform> mesh: Mesh;
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
     [[location(1)]] normal: vec3<f32>;
@@ -21,10 +24,8 @@ struct VertexOutput {
 [[stage(vertex)]]
 fn vertex(vertex: Vertex) -> VertexOutput {
     let position = vertex.position * vertex.i_pos_scale.w + vertex.i_pos_scale.xyz;
-    let world_position = mesh.model * vec4<f32>(position, 1.0);
-
     var out: VertexOutput;
-    out.clip_position = view.view_proj * world_position;
+    out.clip_position = mesh_position_local_to_clip(mesh.model, vec4<f32>(position, 1.0));
     out.color = vertex.i_color;
     return out;
 }

assets/shaders/shader_defs.wgsl

@@ -4,6 +4,9 @@
 [[group(1), binding(0)]]
 var<uniform> mesh: Mesh;
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
     [[location(1)]] normal: vec3<f32>;
@@ -16,10 +19,8 @@ struct VertexOutput {
 
 [[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.clip_position = mesh_position_local_to_clip(mesh.model, vec4<f32>(vertex.position, 1.0));
     return out;
 }
 

crates/bevy_pbr/src/render/depth.wgsl

@@ -13,6 +13,9 @@ var<uniform> joint_matrices: SkinnedMesh;
 #import bevy_pbr::skinning
 #endif
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
 #ifdef SKINNED
@@ -34,6 +37,6 @@ fn vertex(vertex: Vertex) -> VertexOutput {
 #endif
 
     var out: VertexOutput;
-    out.clip_position = view.view_proj * model * vec4<f32>(vertex.position, 1.0);
+    out.clip_position = mesh_position_local_to_clip(model, vec4<f32>(vertex.position, 1.0));
     return out;
 }

crates/bevy_pbr/src/render/mesh.rs

@@ -45,6 +45,8 @@ pub const MESH_TYPES_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 2506024101911992377);
 pub const MESH_BINDINGS_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 16831548636314682308);
+pub const MESH_FUNCTIONS_HANDLE: HandleUntyped =
+    HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 6300874327833745635);
 pub const MESH_SHADER_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 3252377289100772450);
 pub const SKINNING_HANDLE: HandleUntyped =
@@ -71,6 +73,12 @@ impl Plugin for MeshRenderPlugin {
             "mesh_bindings.wgsl",
             Shader::from_wgsl
         );
+        load_internal_asset!(
+            app,
+            MESH_FUNCTIONS_HANDLE,
+            "mesh_functions.wgsl",
+            Shader::from_wgsl
+        );
         load_internal_asset!(app, MESH_SHADER_HANDLE, "mesh.wgsl", Shader::from_wgsl);
         load_internal_asset!(app, SKINNING_HANDLE, "skinning.wgsl", Shader::from_wgsl);
 

crates/bevy_pbr/src/render/mesh.wgsl

@@ -1,6 +1,9 @@
 #import bevy_pbr::mesh_view_bindings
 #import bevy_pbr::mesh_bindings
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
     [[location(1)]] normal: vec3<f32>;
@@ -35,35 +38,21 @@ fn vertex(vertex: Vertex) -> VertexOutput {
     var out: VertexOutput;
 #ifdef SKINNED
     var model = skin_model(vertex.joint_indices, vertex.joint_weights);
-    out.world_position = model * vec4<f32>(vertex.position, 1.0);
     out.world_normal = skin_normals(model, vertex.normal);
-#ifdef VERTEX_TANGENTS
-    out.world_tangent = skin_tangents(model, vertex.tangent);
-#endif
 #else
-    out.world_position = mesh.model * vec4<f32>(vertex.position, 1.0);
-    out.world_normal = mat3x3<f32>(
-        mesh.inverse_transpose_model[0].xyz,
-        mesh.inverse_transpose_model[1].xyz,
-        mesh.inverse_transpose_model[2].xyz
-    ) * vertex.normal;
-#ifdef VERTEX_TANGENTS
-    out.world_tangent = vec4<f32>(
-        mat3x3<f32>(
-            mesh.model[0].xyz,
-            mesh.model[1].xyz,
-            mesh.model[2].xyz
-        ) * vertex.tangent.xyz,
-        vertex.tangent.w
-    );
+    var model = mesh.model;
 #endif
+    out.world_position = mesh_position_local_to_world(model, vec4<f32>(vertex.position, 1.0));
+    out.world_normal = mesh_normal_local_to_world(vertex.normal);
+    out.uv = vertex.uv;
+#ifdef VERTEX_TANGENTS
+    out.world_tangent = mesh_tangent_local_to_world(model, vertex.tangent);
 #endif
 #ifdef VERTEX_COLORS
     out.color = vertex.color;
 #endif
 
-    out.uv = vertex.uv;
-    out.clip_position = view.view_proj * out.world_position;
+    out.clip_position = mesh_position_world_to_clip(out.world_position);
     return out;
 }
 

crates/bevy_pbr/src/render/mesh_functions.wgsl

@@ -0,0 +1,36 @@
+#define_import_path bevy_pbr::mesh_functions
+
+fn mesh_position_local_to_world(model: mat4x4<f32>, vertex_position: vec4<f32>) -> vec4<f32> {
+    return model * vertex_position;
+}
+
+fn mesh_position_world_to_clip(world_position: vec4<f32>) -> vec4<f32> {
+    return view.view_proj * world_position;
+}
+
+// NOTE: The intermediate world_position assignment is important
+// for precision purposes when using the 'equals' depth comparison
+// function.
+fn mesh_position_local_to_clip(model: mat4x4<f32>, vertex_position: vec4<f32>) -> vec4<f32> {
+    let world_position = mesh_position_local_to_world(model, vertex_position);
+    return mesh_position_world_to_clip(world_position);
+}
+
+fn mesh_normal_local_to_world(vertex_normal: vec3<f32>) -> vec3<f32> {
+    return mat3x3<f32>(
+        mesh.inverse_transpose_model[0].xyz,
+        mesh.inverse_transpose_model[1].xyz,
+        mesh.inverse_transpose_model[2].xyz
+    ) * vertex_normal;
+}
+
+fn mesh_tangent_local_to_world(model: mat4x4<f32>, vertex_tangent: vec4<f32>) -> vec4<f32> {
+    return vec4<f32>(
+        mat3x3<f32>(
+            model[0].xyz,
+            model[1].xyz,
+            model[2].xyz
+        ) * vertex_tangent.xyz,
+        vertex_tangent.w
+    );
+}

crates/bevy_pbr/src/render/skinning.wgsl

@@ -1,10 +1,10 @@
-// If using this WGSL snippet as an #import, a dedicated 
+// If using this WGSL snippet as an #import, a dedicated
 // "joint_matricies" uniform of type SkinnedMesh must be added in the
 // main shader.
 
 #define_import_path bevy_pbr::skinning
 
-/// HACK: This works around naga not supporting matrix addition in SPIR-V 
+/// HACK: This works around naga not supporting matrix addition in SPIR-V
 // translations. See https://github.com/gfx-rs/naga/issues/1527
 fn add_matrix(
     a: mat4x4<f32>,
@@ -30,7 +30,7 @@ fn skin_model(
 
 fn inverse_transpose_3x3(in: mat3x3<f32>) -> mat3x3<f32> {
     let x = cross(in.y, in.z);
-    let y = cross(in.z, in.x); 
+    let y = cross(in.z, in.x);
     let z = cross(in.x, in.y);
     let det = dot(in.z, z);
     return mat3x3<f32>(
@@ -50,17 +50,3 @@ fn skin_normals(
         model[2].xyz
     )) * normal;
 }
-
-fn skin_tangents(
-    model: mat4x4<f32>,
-    tangent: vec4<f32>,
-) -> vec4<f32> {
-    return vec4<f32>(
-        mat3x3<f32>(
-            model[0].xyz,
-            model[1].xyz,
-            model[2].xyz
-        ) * tangent.xyz,
-        tangent.w
-    );
-}

crates/bevy_pbr/src/render/wireframe.wgsl

@@ -1,6 +1,18 @@
 #import bevy_pbr::mesh_types
 #import bevy_pbr::mesh_view_bindings
 
+[[group(1), binding(0)]]
+var<uniform> mesh: Mesh;
+
+#ifdef SKINNED
+[[group(1), binding(1)]]
+var<uniform> joint_matrices: SkinnedMesh;
+#import bevy_pbr::skinning
+#endif
+
+// NOTE: Bindings must come before functions that use them!
+#import bevy_pbr::mesh_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
 #ifdef SKINNED
@@ -9,19 +21,10 @@ struct Vertex {
 #endif
 };
 
-[[group(1), binding(0)]]
-var<uniform> mesh: Mesh;
-
 struct VertexOutput {
     [[builtin(position)]] clip_position: vec4<f32>;
 };
 
-#ifdef SKINNED
-[[group(1), binding(1)]]
-var<uniform> joint_matrices: SkinnedMesh;
-#import bevy_pbr::skinning
-#endif
-
 [[stage(vertex)]]
 fn vertex(vertex: Vertex) -> VertexOutput {
 #ifdef SKINNED
@@ -30,10 +33,8 @@ fn vertex(vertex: Vertex) -> VertexOutput {
     let model = mesh.model;
 #endif
 
-    let world_position = model * vec4<f32>(vertex.position, 1.0);
     var out: VertexOutput;
-    out.clip_position = view.view_proj * world_position;
-
+    out.clip_position = mesh_position_local_to_clip(model, vec4<f32>(vertex.position, 1.0));
     return out;
 }
 

crates/bevy_sprite/src/mesh2d/mesh.rs

@@ -45,6 +45,8 @@ pub const MESH2D_TYPES_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 8994673400261890424);
 pub const MESH2D_BINDINGS_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 8983617858458862856);
+pub const MESH2D_FUNCTIONS_HANDLE: HandleUntyped =
+    HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 4976379308250389413);
 pub const MESH2D_SHADER_HANDLE: HandleUntyped =
     HandleUntyped::weak_from_u64(Shader::TYPE_UUID, 2971387252468633715);
 
@@ -74,6 +76,12 @@ impl Plugin for Mesh2dRenderPlugin {
             "mesh2d_bindings.wgsl",
             Shader::from_wgsl
         );
+        load_internal_asset!(
+            app,
+            MESH2D_FUNCTIONS_HANDLE,
+            "mesh2d_functions.wgsl",
+            Shader::from_wgsl
+        );
         load_internal_asset!(app, MESH2D_SHADER_HANDLE, "mesh2d.wgsl", Shader::from_wgsl);
 
         app.add_plugin(UniformComponentPlugin::<Mesh2dUniform>::default());

crates/bevy_sprite/src/mesh2d/mesh2d.wgsl

@@ -1,6 +1,9 @@
 #import bevy_sprite::mesh2d_view_bindings
 #import bevy_sprite::mesh2d_bindings
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_sprite::mesh2d_functions
+
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
     [[location(1)]] normal: vec3<f32>;
@@ -28,26 +31,13 @@ struct VertexOutput {
 
 [[stage(vertex)]]
 fn vertex(vertex: Vertex) -> VertexOutput {
-    let world_position = mesh.model * vec4<f32>(vertex.position, 1.0);
-
     var out: VertexOutput;
     out.uv = vertex.uv;
-    out.world_position = world_position;
-    out.clip_position = view.view_proj * world_position;
-    out.world_normal = mat3x3<f32>(
-        mesh.inverse_transpose_model[0].xyz,
-        mesh.inverse_transpose_model[1].xyz,
-        mesh.inverse_transpose_model[2].xyz
-    ) * vertex.normal;
+    out.world_position = mesh2d_position_local_to_world(mesh.model, vec4<f32>(vertex.position, 1.0));
+    out.clip_position = mesh2d_position_world_to_clip(out.world_position);
+    out.world_normal = mesh2d_normal_local_to_world(vertex.normal);
 #ifdef VERTEX_TANGENTS
-    out.world_tangent = vec4<f32>(
-        mat3x3<f32>(
-            mesh.model[0].xyz,
-            mesh.model[1].xyz,
-            mesh.model[2].xyz
-        ) * vertex.tangent.xyz,
-        vertex.tangent.w
-    );
+    out.world_tangent = mesh2d_tangent_local_to_world(vertex.tangent);
 #endif
 #ifdef VERTEX_COLORS
     out.colors = vertex.colors;

crates/bevy_sprite/src/mesh2d/mesh2d_functions.wgsl

@@ -0,0 +1,36 @@
+#define_import_path bevy_sprite::mesh2d_functions
+
+fn mesh2d_position_local_to_world(model: mat4x4<f32>, vertex_position: vec4<f32>) -> vec4<f32> {
+    return model * vertex_position;
+}
+
+fn mesh2d_position_world_to_clip(world_position: vec4<f32>) -> vec4<f32> {
+    return view.view_proj * world_position;
+}
+
+// NOTE: The intermediate world_position assignment is important
+// for precision purposes when using the 'equals' depth comparison
+// function.
+fn mesh2d_position_local_to_clip(model: mat4x4<f32>, vertex_position: vec4<f32>) -> vec4<f32> {
+    let world_position = mesh2d_position_local_to_world(model, vertex_position);
+    return mesh2d_position_world_to_clip(world_position);
+}
+
+fn mesh2d_normal_local_to_world(vertex_normal: vec3<f32>) -> vec3<f32> {
+    return mat3x3<f32>(
+        mesh.inverse_transpose_model[0].xyz,
+        mesh.inverse_transpose_model[1].xyz,
+        mesh.inverse_transpose_model[2].xyz
+    ) * vertex_normal;
+}
+
+fn mesh2d_tangent_local_to_world(model: mat4x4<f32>, vertex_tangent: vec4<f32>) -> vec4<f32> {
+    return vec4<f32>(
+        mat3x3<f32>(
+            model[0].xyz,
+            model[1].xyz,
+            model[2].xyz
+        ) * vertex_tangent.xyz,
+        vertex_tangent.w
+    );
+}

crates/bevy_sprite/src/render/sprite.wgsl

@@ -28,7 +28,7 @@ fn vertex(
     out.color = vertex_color;
 #endif
     return out;
-} 
+}
 
 [[group(1), binding(0)]]
 var sprite_texture: texture_2d<f32>;
@@ -37,9 +37,9 @@ var sprite_sampler: sampler;
 
 [[stage(fragment)]]
 fn fragment(in: VertexOutput) -> [[location(0)]] vec4<f32> {
-    var color = textureSample(sprite_texture, sprite_sampler, in.uv); 
+    var color = textureSample(sprite_texture, sprite_sampler, in.uv);
 #ifdef COLORED
     color = in.color * color;
 #endif
     return color;
-}
+}

examples/2d/mesh2d_manual.rs

@@ -215,6 +215,9 @@ const COLORED_MESH2D_SHADER: &str = r"
 [[group(1), binding(0)]]
 var<uniform> mesh: Mesh2d;
 
+// NOTE: Bindings must come before functions that use them!
+#import bevy_sprite::mesh2d_functions
+
 // The structure of the vertex buffer is as specified in `specialize()`
 struct Vertex {
     [[location(0)]] position: vec3<f32>;
@@ -233,7 +236,7 @@ struct VertexOutput {
 fn vertex(vertex: Vertex) -> VertexOutput {
     var out: VertexOutput;
     // Project the world position of the mesh into screen position
-    out.clip_position = view.view_proj * mesh.model * vec4<f32>(vertex.position, 1.0);
+    out.clip_position = mesh2d_position_local_to_clip(mesh.model, vec4<f32>(vertex.position, 1.0));
     // Unpack the `u32` from the vertex buffer into the `vec4<f32>` used by the fragment shader
     out.color = vec4<f32>((vec4<u32>(vertex.color) >> vec4<u32>(0u, 8u, 16u, 24u)) & vec4<u32>(255u)) / 255.0;
     return out;