Add option to compress vertex normal tangent to axis angle

Author: beicause

crates/bevy_mesh/src/mesh.rs

@@ -3,7 +3,7 @@ use bevy_derive::EnumVariantMeta;
 use bevy_ecs::resource::Resource;
 use bevy_math::{
     bounding::{Aabb2d, Aabb3d, BoundingVolume},
-    vec2, Vec2, Vec3, Vec3A, Vec3Swizzles,
+    ops, vec2, vec3, vec4, Mat3, Quat, Vec2, Vec3, Vec3A, Vec3Swizzles, Vec4, Vec4Swizzles,
 };
 #[cfg(feature = "serialize")]
 use bevy_platform::collections::HashMap;
@@ -1155,8 +1155,10 @@ pub fn octahedral_encode_signed(v: Vec3) -> Vec2 {
 
 /// Encode tangent vectors as octahedral coordinates with range [-1, 1]. The sign is encoded in y component. Use [`octahedral_decode_tangent`] to decode.
 pub fn octahedral_encode_tangent(v: Vec3, sign: f32) -> Vec2 {
-    // Bias to ensure that encoding as unorm16 preserves the sign. See https://github.com/godotengine/godot/pull/73265
-    let bias = 1.0 / 32767.0;
+    // Bias to ensure that encoding as snorm16 preserves the sign.
+    let bits = 16.;
+    let bias = 1. / (ops::powf(2.0, bits - 1.) - 1.);
+
     let mut n_xy = octahedral_encode_signed(v);
     // Map y to always be positive.
     n_xy.y = n_xy.y * 0.5 + 0.5;
@@ -1184,12 +1186,60 @@ pub fn octahedral_decode_tangent(v: Vec2) -> (Vec3, f32) {
     (octahedral_decode_signed(f), sign)
 }
 
+/// Convert the normal and tangent to the equivalent axis-angle representation.
+/// The range of angle is [-2pi, 2pi], where the sign represents the handedness of the tangent.
+pub fn normal_tangent_to_axis_angle(normal: Vec3, tangent_signed: Vec4) -> (Vec3, f32) {
+    // Bias to ensure that encoding as snorm16 preserves the sign.
+    let bits = 16.;
+    let bias = 1. / (ops::powf(2.0, bits - 1.) - 1.);
+
+    let tangent = tangent_signed.xyz();
+    let bitangent = normal.cross(tangent);
+    let (axis, angle) =
+        Quat::from_mat3(&Mat3::from_cols(tangent, bitangent, normal)).to_axis_angle();
+    let angle = angle
+        .rem_euclid(2.0 * core::f32::consts::PI)
+        .max(bias * 2.0 * core::f32::consts::PI);
+
+    (axis, angle * tangent_signed.w)
+}
+
+/// Convert the axis-angle representation back to normal and tangent.
+/// The range of angle is [-2pi, 2pi], where the sign represents the handedness of the tangent.
+pub fn axis_angle_to_normal_tangent(axis: Vec3, angle: f32) -> (Vec3, Vec4) {
+    let sign = if angle >= 0.0 { 1.0 } else { -1.0 };
+    // References the source code of `Mat3::from_quat(Quat::from_axis_angle(axis, angle))`
+    let angle = angle * 0.5 * sign;
+    let c = ops::cos(angle);
+    let s = ops::sin(angle);
+    let v = axis * s;
+    let rotation = vec4(v.x, v.y, v.z, c);
+    let x2 = rotation.x + rotation.x;
+    let y2 = rotation.y + rotation.y;
+    let z2 = rotation.z + rotation.z;
+    let xx = rotation.x * x2;
+    let xy = rotation.x * y2;
+    let xz = rotation.x * z2;
+    let yy = rotation.y * y2;
+    let yz = rotation.y * z2;
+    let zz = rotation.z * z2;
+    let wx = rotation.w * x2;
+    let wy = rotation.w * y2;
+    let wz = rotation.w * z2;
+
+    let tangent = vec3(1.0 - (yy + zz), xy + wz, xz - wy);
+    let normal = vec3(xz + wy, yz - wx, 1.0 - (xx + yy));
+
+    (normal, tangent.extend(sign))
+}
+
 #[cfg(test)]
 mod tests {
-    use bevy_math::{vec2, vec3, Vec4Swizzles};
+    use bevy_math::{vec2, vec3, Mat3, Quat, Vec3, Vec4, Vec4Swizzles};
 
     use crate::{
-        octahedral_decode_signed, octahedral_decode_tangent,
+        axis_angle_to_normal_tangent, normal_tangent_to_axis_angle, octahedral_decode_signed,
+        octahedral_decode_tangent,
         vertex::{octahedral_encode_signed, octahedral_encode_tangent},
     };
 
@@ -1198,18 +1248,21 @@ mod tests {
         let vectors = [
             vec3(1.0, 2.0, 3.0).normalize().extend(1.0),
             vec3(1.0, 0.0, 0.0).extend(-1.0),
+            vec3(0.0, 1.0, 0.0).extend(-1.0),
             vec3(0.0, 0.0, -1.0).extend(1.0),
             vec3(0.0, 0.0, -1.0).extend(-1.0),
         ];
         let expected_encoded_normals = [
             vec2(0.16666667, 0.33333334),
             vec2(1.0, 0.0),
+            vec2(0.0, 1.0),
             vec2(-1.0, -1.0),
             vec2(-1.0, -1.0),
         ];
         let expected_encoded_tangents = [
             vec2(0.16666667, 0.6666667),
             vec2(1.0, -0.5),
+            vec2(0.0, -1.0),
             vec2(-1.0, 3.051851e-5),
             vec2(-1.0, -3.051851e-5),
         ];
@@ -1226,4 +1279,75 @@ mod tests {
             assert!(decoded_tangent.distance(v.xyz()) < 1e-4);
         }
     }
+
+    pub fn axis_angle_to_normal_tangent_glam(axis: Vec3, angle: f32) -> (Vec3, Vec4) {
+        let sign = if angle >= 0.0 { 1.0 } else { -1.0 };
+        let tbn = Mat3::from_quat(Quat::from_axis_angle(axis, angle * sign));
+        (tbn.col(2), tbn.col(0).extend(sign))
+    }
+
+    #[test]
+    fn normal_tangent_axis_angle_encode_decode() {
+        let normal_tangent = [
+            (vec3(1.0, 0.0, 0.0), vec3(0.0, 0.0, 1.0).extend(1.0)),
+            (vec3(1.0, 0.0, 0.0), vec3(0.0, 0.0, 1.0).extend(-1.0)),
+            (vec3(0.0, 0.0, 1.0), vec3(1.0, 0.0, 0.0).extend(1.0)),
+            (vec3(0.0, 0.0, 1.0), vec3(1.0, 0.0, 0.0).extend(-1.0)),
+            (vec3(0.0, 1.0, 0.0), vec3(0.0, 0.0, 1.0).extend(1.0)),
+            (vec3(1.0, 0.0, 0.0), vec3(0.0, -1.0, 0.0).extend(1.0)),
+            (
+                vec3(1.0, 1.0, 1.0).normalize(),
+                vec3(1.0, -1.0, 0.0).normalize().extend(1.0),
+            ),
+            (
+                vec3(1.0, 2.0, 0.0).normalize(),
+                vec3(0.0, 0.0, 1.0).extend(-1.0),
+            ),
+            (
+                vec3(0.0, 1.0, 1.0).normalize(),
+                vec3(1.0, 0.0, 0.0).extend(1.0),
+            ),
+            (
+                vec3(-1.0, 1.0, 1.0).normalize(),
+                vec3(1.0, 1.0, 0.0).normalize().extend(-1.0),
+            ),
+            (
+                vec3(3.0, 1.0, 2.0).normalize(),
+                vec3(0.0, 1.0, -0.5).normalize().extend(1.0),
+            ),
+        ];
+
+        #[expect(
+            clippy::approx_constant,
+            reason = "The values are taken from the test results"
+        )]
+        let expected_axis_angle = [
+            (vec3(0.7071068, 0.0, 0.7071068), 3.1415927),
+            (vec3(0.7071068, 0.0, 0.7071068), -3.1415927),
+            (vec3(1.0, 0.0, 0.0), 3.051851e-5),
+            (vec3(1.0, 0.0, 0.0), -3.051851e-5),
+            (vec3(0.57735026, 0.57735026, 0.57735026), 4.1887903),
+            (vec3(0.57735026, -0.57735026, 0.57735026), 4.1887903),
+            (vec3(-0.7429061, 0.3077218, -0.5944728), 1.2171159),
+            (vec3(0.64793617, 0.40044653, 0.64793617), -3.9033751),
+            (vec3(-1.0, 0.0, 0.0), 0.7853981),
+            (vec3(-0.7429061, -0.3077218, 0.5944728), -1.2171159),
+            (vec3(0.22525999, 0.6253928, 0.7470889), 1.6242763),
+        ];
+
+        for (i, &(normal, tangent)) in normal_tangent.iter().enumerate() {
+            let (axis, angle) = normal_tangent_to_axis_angle(normal, tangent);
+            assert_eq!(angle.signum(), tangent.w.signum());
+            assert!(axis.distance(expected_axis_angle[i].0) < 1e-8);
+
+            let (decoded_normal, decoded_tangent) = axis_angle_to_normal_tangent(axis, angle);
+            let (decoded_normal_glam, decoded_tangent_glam) =
+                axis_angle_to_normal_tangent_glam(axis, angle);
+
+            assert!(decoded_normal.distance(normal) < 1e-3);
+            assert!(decoded_tangent.distance(tangent) < 1e-3);
+            assert!(decoded_normal_glam.distance(normal) < 1e-3);
+            assert!(decoded_tangent_glam.distance(tangent) < 1e-3);
+        }
+    }
 }

crates/bevy_mesh/src/vertex.rs

@@ -535,6 +535,14 @@ impl PrepassPipeline {
                 }
                 vertex_attributes.push(Mesh::ATTRIBUTE_TANGENT.at_shader_location(4));
             }
+            if layout
+                .0
+                .get_attribute_compression()
+                .contains(MeshAttributeCompressionFlags::PACKED_AXIS_ANGLE_TBN)
+            {
+                shader_defs.push("VERTEX_TANGENTS".into());
+                shader_defs.push("VERTEX_PACKED_AXIS_ANGLE_TBN".into());
+            }
         }
         if mesh_key
             .intersects(MeshPipelineKey::MOTION_VECTOR_PREPASS | MeshPipelineKey::DEFERRED_PREPASS)

crates/bevy_pbr/src/prepass/mod.rs

@@ -18,9 +18,13 @@ struct UncompressedVertex {
 #ifdef VERTEX_NORMALS
     @location(3) normal: vec3<f32>,
 #endif
+
 #ifdef VERTEX_TANGENTS
     @location(4) tangent: vec4<f32>,
+#else ifdef VERTEX_PACKED_AXIS_ANGLE_TBN
+    @location(4) tangent: vec4<f32>,
 #endif
+
 #endif // NORMAL_PREPASS_OR_DEFERRED_PREPASS
 
 #ifdef SKINNED
@@ -39,26 +43,20 @@ struct UncompressedVertex {
 
 struct Vertex {
     @builtin(instance_index) instance_index: u32,
+
+#ifdef VERTEX_PACKED_AXIS_ANGLE_TBN
+    @location(0) compressed_position_angle: vec4<f32>,
+#ifdef NORMAL_PREPASS_OR_DEFERRED_PREPASS
+    @location(3) compressed_axis: vec2<f32>,
+#endif
+
+#else // VERTEX_PACKED_AXIS_ANGLE_TBN
+
 #ifdef VERTEX_POSITIONS_COMPRESSED
     @location(0) compressed_position: vec4<f32>,
 #else
     @location(0) position: vec3<f32>,
 #endif
-#ifdef VERTEX_UVS_A
-#ifdef VERTEX_UVS_A_COMPRESSED
-    @location(1) compressed_uv: vec2<f32>,
-#else
-    @location(1) uv: vec2<f32>,
-#endif
-#endif
-#ifdef VERTEX_UVS_B
-#ifdef VERTEX_UVS_B_COMPRESSED
-    @location(2) compressed_uv_b: vec2<f32>,
-#else
-    @location(2) uv_b: vec2<f32>,
-#endif
-#endif
-
 #ifdef NORMAL_PREPASS_OR_DEFERRED_PREPASS
 #ifdef VERTEX_NORMALS
 #ifdef VERTEX_NORMALS_COMPRESSED
@@ -76,6 +74,23 @@ struct Vertex {
 #endif
 #endif // NORMAL_PREPASS_OR_DEFERRED_PREPASS
 
+#endif // VERTEX_PACKED_AXIS_ANGLE_TBN
+
+#ifdef VERTEX_UVS_A
+#ifdef VERTEX_UVS_A_COMPRESSED
+    @location(1) compressed_uv: vec2<f32>,
+#else
+    @location(1) uv: vec2<f32>,
+#endif
+#endif
+#ifdef VERTEX_UVS_B
+#ifdef VERTEX_UVS_B_COMPRESSED
+    @location(2) compressed_uv_b: vec2<f32>,
+#else
+    @location(2) uv_b: vec2<f32>,
+#endif
+#endif
+
 #ifdef SKINNED
     @location(5) joint_indices: vec4<u32>,
     @location(6) joint_weights: vec4<f32>,
@@ -96,6 +111,24 @@ fn decompress_vertex(vertex_in: Vertex, instance_index: u32) -> UncompressedVert
     let mesh_metadata = bevy_pbr::mesh_functions::get_metadata(instance_index);
     var uncompressed_vertex: UncompressedVertex;
     uncompressed_vertex.instance_index = instance_index;
+
+#ifdef VERTEX_PACKED_AXIS_ANGLE_TBN
+    uncompressed_vertex.position = bevy_render::utils::decompress_vertex_position(vertex_in.compressed_position_angle, mesh_metadata.aabb_center, mesh_metadata.aabb_half_extents);
+#ifdef NORMAL_PREPASS_OR_DEFERRED_PREPASS
+    var normal: vec3f;
+    var tangent: vec4f;
+    bevy_render::utils::decompress_vertex_axis_angle_to_normal_tangent(
+        vertex_in.compressed_axis,
+        vertex_in.compressed_position_angle.w,
+        &normal,
+        &tangent,
+    );
+    uncompressed_vertex.normal = normal;
+    uncompressed_vertex.tangent = tangent;
+#endif // NORMAL_PREPASS_OR_DEFERRED_PREPASS
+
+#else // VERTEX_PACKED_AXIS_ANGLE_TBN
+
 #ifdef VERTEX_POSITIONS_COMPRESSED
     uncompressed_vertex.position = bevy_render::utils::decompress_vertex_position(vertex_in.compressed_position, mesh_metadata.aabb_center, mesh_metadata.aabb_half_extents);
 #else
@@ -117,6 +150,9 @@ fn decompress_vertex(vertex_in: Vertex, instance_index: u32) -> UncompressedVert
 #endif
 #endif
 #endif // NORMAL_PREPASS_OR_DEFERRED_PREPASS
+
+#endif // VERTEX_PACKED_AXIS_ANGLE_TBN
+
 #ifdef VERTEX_UVS_A
 #ifdef VERTEX_UVS_A_COMPRESSED
     let uv_min_and_extents_a = mesh_metadata.uv_channels_min_and_extents[0];