Fast renormalize (#14316)

# Objective

- Addresses part of #14302 .

## Solution

- Add a fast_remormalize method to Dir2/Dir3/Dir3A and Rot2.

## Testing

- Added tests too

---------

Co-authored-by: BD103 <59022059+BD103@users.noreply.github.com>
Co-authored-by: Jan Hohenheim <jan@hohenheim.ch>
Co-authored-by: Joona Aalto <jondolf.dev@gmail.com>

crates/bevy_math/src/direction.rs

@@ -243,6 +243,16 @@ impl Dir2 {
     pub fn rotation_to_y(self) -> Rot2 {
         self.rotation_from_y().inverse()
     }
+
+    /// Returns `self` after an approximate normalization, assuming the value is already nearly normalized.
+    /// Useful for preventing numerical error accumulation.
+    /// See [`Dir3::fast_renormalize`] for an example of when such error accumulation might occur.
+    #[inline]
+    pub fn fast_renormalize(self) -> Self {
+        let length_squared = self.0.length_squared();
+        // Based on a Taylor approximation of the inverse square root, see [`Dir3::fast_renormalize`] for more details.
+        Self(self * (0.5 * (3.0 - length_squared)))
+    }
 }
 
 impl TryFrom<Vec2> for Dir2 {
@@ -446,6 +456,56 @@ impl Dir3 {
         let quat = Quat::IDENTITY.slerp(Quat::from_rotation_arc(self.0, rhs.0), s);
         Dir3(quat.mul_vec3(self.0))
     }
+
+    /// Returns `self` after an approximate normalization, assuming the value is already nearly normalized.
+    /// Useful for preventing numerical error accumulation.
+    ///
+    /// # Example
+    /// The following seemingly benign code would start accumulating errors over time,
+    /// leading to `dir` eventually not being normalized anymore.
+    /// ```
+    /// # use bevy_math::prelude::*;
+    /// # let N: usize = 200;
+    /// let mut dir = Dir3::X;
+    /// let quaternion = Quat::from_euler(EulerRot::XYZ, 1.0, 2.0, 3.0);
+    /// for i in 0..N {
+    ///     dir = quaternion * dir;
+    /// }
+    /// ```
+    /// Instead, do the following.
+    /// ```
+    /// # use bevy_math::prelude::*;
+    /// # let N: usize = 200;
+    /// let mut dir = Dir3::X;
+    /// let quaternion = Quat::from_euler(EulerRot::XYZ, 1.0, 2.0, 3.0);
+    /// for i in 0..N {
+    ///     dir = quaternion * dir;
+    ///     dir = dir.fast_renormalize();
+    /// }
+    /// ```
+    #[inline]
+    pub fn fast_renormalize(self) -> Self {
+        // We numerically approximate the inverse square root by a Taylor series around 1
+        // As we expect the error (x := length_squared - 1) to be small
+        // inverse_sqrt(length_squared) = (1 + x)^(-1/2) = 1 - 1/2 x + O(x²)
+        // inverse_sqrt(length_squared) ≈ 1 - 1/2 (length_squared - 1) = 1/2 (3 - length_squared)
+
+        // Iterative calls to this method quickly converge to a normalized value,
+        // so long as the denormalization is not large ~ O(1/10).
+        // One iteration can be described as:
+        // l_sq <- l_sq * (1 - 1/2 (l_sq - 1))²;
+        // Rewriting in terms of the error x:
+        // 1 + x <- (1 + x) * (1 - 1/2 x)²
+        // 1 + x <- (1 + x) * (1 - x + 1/4 x²)
+        // 1 + x <- 1 - x + 1/4 x² + x - x² + 1/4 x³
+        // x <- -1/4 x² (3 - x)
+        // If the error is small, say in a range of (-1/2, 1/2), then:
+        // |-1/4 x² (3 - x)| <= (3/4 + 1/4 * |x|) * x² <= (3/4 + 1/4 * 1/2) * x² < x² < 1/2 x
+        // Therefore the sequence of iterates converges to 0 error as a second order method.
+
+        let length_squared = self.0.length_squared();
+        Self(self * (0.5 * (3.0 - length_squared)))
+    }
 }
 
 impl TryFrom<Vec3> for Dir3 {
@@ -655,6 +715,17 @@ impl Dir3A {
         );
         Dir3A(quat.mul_vec3a(self.0))
     }
+
+    /// Returns `self` after an approximate normalization, assuming the value is already nearly normalized.
+    /// Useful for preventing numerical error accumulation.
+    ///
+    /// See [`Dir3::fast_renormalize`] for an example of when such error accumulation might occur.
+    #[inline]
+    pub fn fast_renormalize(self) -> Self {
+        let length_squared = self.0.length_squared();
+        // Based on a Taylor approximation of the inverse square root, see [`Dir3::fast_renormalize`] for more details.
+        Self(self * (0.5 * (3.0 - length_squared)))
+    }
 }
 
 impl From<Dir3> for Dir3A {
@@ -815,6 +886,31 @@ mod tests {
         assert_relative_eq!(Dir2::NORTH_WEST.rotation_from_y(), Rot2::FRAC_PI_4);
     }
 
+    #[test]
+    fn dir2_renorm() {
+        // Evil denormalized Rot2
+        let (sin, cos) = 1.0_f32.sin_cos();
+        let rot2 = Rot2::from_sin_cos(sin * (1.0 + 1e-5), cos * (1.0 + 1e-5));
+        let mut dir_a = Dir2::X;
+        let mut dir_b = Dir2::X;
+
+        // We test that renormalizing an already normalized dir doesn't do anything
+        assert_relative_eq!(dir_b, dir_b.fast_renormalize(), epsilon = 0.000001);
+
+        for _ in 0..50 {
+            dir_a = rot2 * dir_a;
+            dir_b = rot2 * dir_b;
+            dir_b = dir_b.fast_renormalize();
+        }
+
+        // `dir_a` should've gotten denormalized, meanwhile `dir_b` should stay normalized.
+        assert!(
+            !dir_a.is_normalized(),
+            "Dernormalization doesn't work, test is faulty"
+        );
+        assert!(dir_b.is_normalized(), "Renormalisation did not work.");
+    }
+
     #[test]
     fn dir3_creation() {
         assert_eq!(Dir3::new(Vec3::X * 12.5), Ok(Dir3::X));
@@ -862,6 +958,30 @@ mod tests {
         );
     }
 
+    #[test]
+    fn dir3_renorm() {
+        // Evil denormalized quaternion
+        let rot3 = Quat::from_euler(glam::EulerRot::XYZ, 1.0, 2.0, 3.0) * (1.0 + 1e-5);
+        let mut dir_a = Dir3::X;
+        let mut dir_b = Dir3::X;
+
+        // We test that renormalizing an already normalized dir doesn't do anything
+        assert_relative_eq!(dir_b, dir_b.fast_renormalize(), epsilon = 0.000001);
+
+        for _ in 0..50 {
+            dir_a = rot3 * dir_a;
+            dir_b = rot3 * dir_b;
+            dir_b = dir_b.fast_renormalize();
+        }
+
+        // `dir_a` should've gotten denormalized, meanwhile `dir_b` should stay normalized.
+        assert!(
+            !dir_a.is_normalized(),
+            "Dernormalization doesn't work, test is faulty"
+        );
+        assert!(dir_b.is_normalized(), "Renormalisation did not work.");
+    }
+
     #[test]
     fn dir3a_creation() {
         assert_eq!(Dir3A::new(Vec3A::X * 12.5), Ok(Dir3A::X));
@@ -908,4 +1028,28 @@ mod tests {
             Dir3A::from_xyz(0.0, 0.75f32.sqrt(), 0.5).unwrap()
         );
     }
+
+    #[test]
+    fn dir3a_renorm() {
+        // Evil denormalized quaternion
+        let rot3 = Quat::from_euler(glam::EulerRot::XYZ, 1.0, 2.0, 3.0) * (1.0 + 1e-5);
+        let mut dir_a = Dir3A::X;
+        let mut dir_b = Dir3A::X;
+
+        // We test that renormalizing an already normalized dir doesn't do anything
+        assert_relative_eq!(dir_b, dir_b.fast_renormalize(), epsilon = 0.000001);
+
+        for _ in 0..50 {
+            dir_a = rot3 * dir_a;
+            dir_b = rot3 * dir_b;
+            dir_b = dir_b.fast_renormalize();
+        }
+
+        // `dir_a` should've gotten denormalized, meanwhile `dir_b` should stay normalized.
+        assert!(
+            !dir_a.is_normalized(),
+            "Dernormalization doesn't work, test is faulty"
+        );
+        assert!(dir_b.is_normalized(), "Renormalisation did not work.");
+    }
 }

crates/bevy_math/src/rotation2d.rs

@@ -226,6 +226,20 @@ impl Rot2 {
         Self::from_sin_cos(self.sin * length_recip, self.cos * length_recip)
     }
 
+    /// Returns `self` after an approximate normalization, assuming the value is already nearly normalized.
+    /// Useful for preventing numerical error accumulation.
+    /// See [`Dir3::fast_renormalize`](crate::Dir3::fast_renormalize) for an example of when such error accumulation might occur.
+    #[inline]
+    pub fn fast_renormalize(self) -> Self {
+        let length_squared = self.length_squared();
+        // Based on a Taylor approximation of the inverse square root, see [`Dir3::fast_renormalize`](crate::Dir3::fast_renormalize) for more details.
+        let length_recip_approx = 0.5 * (3.0 - length_squared);
+        Rot2 {
+            sin: self.sin * length_recip_approx,
+            cos: self.cos * length_recip_approx,
+        }
+    }
+
     /// Returns `true` if the rotation is neither infinite nor NaN.
     #[inline]
     pub fn is_finite(self) -> bool {
@@ -522,6 +536,45 @@ mod tests {
         assert!(normalized_rotation.is_normalized());
     }
 
+    #[test]
+    fn fast_renormalize() {
+        let rotation = Rot2 { sin: 1.0, cos: 0.5 };
+        let normalized_rotation = rotation.normalize();
+
+        let mut unnormalized_rot = rotation;
+        let mut renormalized_rot = rotation;
+        let mut initially_normalized_rot = normalized_rotation;
+        let mut fully_normalized_rot = normalized_rotation;
+
+        // Compute a 64x (=2⁶) multiple of the rotation.
+        for _ in 0..6 {
+            unnormalized_rot = unnormalized_rot * unnormalized_rot;
+            renormalized_rot = renormalized_rot * renormalized_rot;
+            initially_normalized_rot = initially_normalized_rot * initially_normalized_rot;
+            fully_normalized_rot = fully_normalized_rot * fully_normalized_rot;
+
+            renormalized_rot = renormalized_rot.fast_renormalize();
+            fully_normalized_rot = fully_normalized_rot.normalize();
+        }
+
+        assert!(!unnormalized_rot.is_normalized());
+
+        assert!(renormalized_rot.is_normalized());
+        assert!(fully_normalized_rot.is_normalized());
+
+        assert_relative_eq!(fully_normalized_rot, renormalized_rot, epsilon = 0.000001);
+        assert_relative_eq!(
+            fully_normalized_rot,
+            unnormalized_rot.normalize(),
+            epsilon = 0.000001
+        );
+        assert_relative_eq!(
+            fully_normalized_rot,
+            initially_normalized_rot.normalize(),
+            epsilon = 0.000001
+        );
+    }
+
     #[test]
     fn try_normalize() {
         // Valid