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

Author: superdump

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,28 +1,29 @@
 #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>;
 };
 
 [[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;
 }