Account for over/underflow on halfAngle and axis calculations

Author: markuswntr

Sources/QuaternionModule/Transformation.swift

@@ -58,10 +58,14 @@ extension Quaternion {
   /// [wiki]: https://en.wikipedia.org/wiki/Quaternions_and_spatial_rotation#Recovering_the_axis-angle_representation
   @inlinable
   public var axis: SIMD3<RealType> {
-    guard isFinite, imaginary != .zero, real != .zero else {
-      return SIMD3(repeating: .nan)
-    }
-    return imaginary / .sqrt(imaginary.lengthSquared)
+    guard isFinite, imaginary != .zero else { return SIMD3(repeating: .nan) }
+
+    // If lengthSquared computes without over/underflow, everything is fine
+    // and the result is correct. If not, we have to do the computation
+    // carefully and unscale the quaternion first.
+    let lenSq = imaginary.lengthSquared
+    guard lenSq.isNormal else { return divided(by: magnitude).axis }
+    return imaginary / .sqrt(lenSq)
   }
 
   /// The [Angle-Axis][wiki] representation.
@@ -267,7 +271,7 @@ extension Quaternion {
   @inlinable
   public init(rotation vector: SIMD3<RealType>) {
     let angle: RealType = .sqrt(vector.lengthSquared)
-    if !angle.isZero && angle.isFinite {
+    if !angle.isZero, angle.isFinite {
       self = Quaternion(halfAngle: angle/2, unitAxis: vector/angle)
     } else {
       self = Quaternion(angle)
@@ -404,8 +408,19 @@ extension Quaternion {
   /// If the quaternion is zero or non-finite, halfAngle is `nan`.
   @usableFromInline @inline(__always)
   internal var halfAngle: RealType {
-    guard !isZero && isFinite else { return .nan }
-    return .atan2(y: .sqrt(imaginary.lengthSquared), x: real)
+    guard isFinite else { return .nan }
+    guard imaginary != .zero else {
+      // A zero quaternion does not encode transformation properties.
+      // If imaginary is zero, real must be non-zero or nan is returned.
+      return real.isZero ? .nan : .zero
+    }
+
+    // If lengthSquared computes without over/underflow, everything is fine
+    // and the result is correct. If not, we have to do the computation
+    // carefully and unscale the quaternion first.
+    let lenSq = imaginary.lengthSquared
+    guard lenSq.isNormal else { return divided(by: magnitude).halfAngle }
+    return .atan2(y: .sqrt(lenSq), x: real)
   }
 
   /// Creates a new quaternion from given half rotation angle about given
@@ -430,18 +445,18 @@ extension Quaternion {
 // and *(x,y,z)* axis representations internally to the module.
 extension SIMD3 where Scalar: FloatingPoint {
 
-  /// Returns the squared length of this instance.
-  @usableFromInline @inline(__always)
-  internal var lengthSquared: Scalar {
-    dot(self)
-  }
-
   /// True if all values of this instance are finite
   @usableFromInline @inline(__always)
   internal var isFinite: Bool {
     x.isFinite && y.isFinite && z.isFinite
   }
 
+  /// Returns the squared length of this instance.
+  @usableFromInline @inline(__always)
+  internal var lengthSquared: Scalar {
+    dot(self)
+  }
+
   /// Returns the scalar/dot product of this vector with `other`.
   @usableFromInline @inline(__always)
   internal func dot(_ other: SIMD3<Scalar>) -> Scalar {

Tests/QuaternionTests/TransformationTests.swift

@@ -116,6 +116,36 @@ final class TransformationTests: XCTestCase {
     testAngleAxisEdgeCases(Float64.self)
   }
 
+  func testHalfAngleAndAxisOverflow<T: Real & SIMDScalar>(_ type: T.Type) {
+    let unscaled = Quaternion<T>(1, SIMD3(repeating: 1))
+    let scaled = Quaternion<T>(
+      .greatestFiniteMagnitude,
+      SIMD3(repeating: .greatestFiniteMagnitude)
+    )
+    XCTAssertEqual(scaled.angle, unscaled.angle)
+    XCTAssertEqual(scaled.axis, unscaled.axis)
+  }
+
+  func testHalfAngleAndAxisOverflow() {
+    testHalfAngleAndAxisOverflow(Float32.self)
+    testHalfAngleAndAxisOverflow(Float64.self)
+  }
+
+  func testHalfAngleAndAxisUnderflow<T: Real & SIMDScalar>(_ type: T.Type) {
+    let unscaled = Quaternion<T>(1, SIMD3(repeating: 1))
+    let scaled = Quaternion<T>(
+      .leastNormalMagnitude,
+      SIMD3(repeating: .leastNormalMagnitude)
+    )
+    XCTAssertEqual(scaled.angle, unscaled.angle)
+    XCTAssertEqual(scaled.axis, unscaled.axis)
+  }
+
+  func testHalfAngleAndAxisUnderflow() {
+    testHalfAngleAndAxisUnderflow(Float32.self)
+    testHalfAngleAndAxisUnderflow(Float64.self)
+  }
+
   // MARK: Rotation Vector
 
   func testRotationVector<T: Real & SIMDScalar>(_ type: T.Type) {
@@ -170,11 +200,11 @@ final class TransformationTests: XCTestCase {
     XCTAssertEqual(Quaternion<T>(length: .zero, phase: .infinity, axis:  .infinity), .zero)
     XCTAssertEqual(Quaternion<T>(length: .zero, phase:-.infinity, axis: -.infinity), .zero)
     XCTAssertEqual(Quaternion<T>(length: .zero, phase: .nan     , axis:  .nan     ), .zero)
-    XCTAssertEqual(Quaternion<T>(length: .infinity, phase: .zero    , axis:  .zero    ), .zero)
+    XCTAssertEqual(Quaternion<T>(length: .infinity, phase: .zero    , axis:  .zero    ), .infinity)
     XCTAssertEqual(Quaternion<T>(length: .infinity, phase: .infinity, axis:  .infinity), .infinity)
     XCTAssertEqual(Quaternion<T>(length: .infinity, phase:-.infinity, axis: -.infinity), .infinity)
     XCTAssertEqual(Quaternion<T>(length: .infinity, phase: .nan     , axis:  .infinity), .infinity)
-    XCTAssertEqual(Quaternion<T>(length:-.infinity, phase: .zero    , axis:  .zero    ), .zero)
+    XCTAssertEqual(Quaternion<T>(length:-.infinity, phase: .zero    , axis:  .zero    ), .infinity)
     XCTAssertEqual(Quaternion<T>(length:-.infinity, phase: .infinity, axis:  .infinity), .infinity)
     XCTAssertEqual(Quaternion<T>(length:-.infinity, phase:-.infinity, axis: -.infinity), .infinity)
     XCTAssertEqual(Quaternion<T>(length:-.infinity, phase: .nan     , axis:  .infinity), .infinity)