Test coverage for saturating shift + fix the bug that it turned up.

Author: stephentyrone

Sources/IntegerUtilities/SaturatingArithmetic.swift

@@ -11,7 +11,7 @@
 
 extension FixedWidthInteger {
   @_transparent @usableFromInline
-  var signExtension: Self { self &>> -1 }
+  var sextOrZext: Self { self >> Self.bitWidth }
 
   /// Saturating integer addition
   ///
@@ -30,7 +30,7 @@ extension FixedWidthInteger {
   public func addingWithSaturation(_ other: Self) -> Self {
     let (wrapped, overflow) = addingReportingOverflow(other)
     if !overflow { return wrapped }
-    return Self.max &- signExtension
+    return Self.max &- sextOrZext
   }
 
   /// Saturating integer subtraction
@@ -54,7 +54,7 @@ extension FixedWidthInteger {
   public func subtractingWithSaturation(_ other: Self) -> Self {
     let (wrapped, overflow) = subtractingReportingOverflow(other)
     if !overflow { return wrapped }
-    return Self.max &- signExtension
+    return Self.max &- sextOrZext
   }
 
   /// Saturating integer negation
@@ -85,8 +85,8 @@ extension FixedWidthInteger {
   public func multipliedWithSaturation(by other: Self) -> Self {
     let (high, low) = multipliedFullWidth(by: other)
     let wrapped = Self(truncatingIfNeeded: low)
-    if high == wrapped.signExtension { return wrapped }
-    return Self.max &- high.signExtension
+    if high == wrapped.sextOrZext { return wrapped }
+    return Self.max &- high.sextOrZext
   }
 
   /// Bitwise left with rounding and saturation.
@@ -106,23 +106,45 @@ extension FixedWidthInteger {
     leftBy count: Count, rounding rule: RoundingRule = .down
   ) -> Self {
     // If count is zero or negative, negate it and do a right
-    // shift without saturation instead, as that's easier.
+    // shift without saturation instead, since we already have
+    // that implemented.
     guard count > 0 else {
+      // negating count is tricky, because count's type can be
+      // an arbitrary BinaryInteger; in particular, it could be
+      // .min of a signed type, so that its negation cannot be
+      // represented in the same type. Fortunately, Int64 is
+      // always big enough to represent arbitrary shifts of
+      // arbitrary types, so we can use that as an intermediate
+      // type, and then we can use negatedWithSaturation() to
+      // handle the .min case.
+      let int64Count = Int64(clamping: count)
       return shifted(
-        rightBy: Self(clamping: count).negatedWithSaturation(),
+        rightBy: int64Count.negatedWithSaturation(),
         rounding: rule
       )
     }
+    let clamped = Self.max &- sextOrZext
     guard count < Self.bitWidth else {
       // If count is bitWidth or greater, we always overflow
       // unless self is zero.
-      return self == 0 ? 0 : Self.max &- signExtension
+      return self == 0 ? 0 : clamped
     }
     // Now we have 0 < count < bitWidth, so we can use a nice
-    // straightforward implementation; the shift overflows if
-    // the complementary shift doesn't match sign extension.
-    let wrapped = self << count
-    if self &>> ~count == signExtension { return wrapped }
-    return Self.max &- signExtension
+    // straightforward implementation; a shift overflows if
+    // the complementary shift doesn't match sign-or-zero
+    // extension. E.g.:
+    //
+    // - signed 0b0010_1111 << 2 overflows, because
+    //   0b0010_1111 >> 5 is 0b0000_0001, which does not
+    //   equal 0b0000_0000
+    //
+    // - unsigned 0b0010_1111 << 2 does not overflow,
+    //   because 0b0010_0000 >> 6 is 0b0000_0000, which
+    //   does equal 0b0000_0000.
+    let valueBits = Self.bitWidth &- (Self.isSigned ? 1 : 0)
+    let wrapped = self &<< count
+    let complement = valueBits &- Int(count)
+    if self &>> complement == sextOrZext { return wrapped }
+    return clamped
   }
 }

Tests/IntegerUtilitiesTests/ShiftTests.swift

@@ -16,9 +16,9 @@ import _TestSupport
 
 final class IntegerUtilitiesShiftTests: XCTestCase {
 
-  func testRoundingShift<T: FixedWidthInteger>(
-    _ value: T, _ count: Int, rounding rule: RoundingRule
-  ) {
+  func testRoundingShift<T, C>(
+    _ value: T, _ count: C, rounding rule: RoundingRule
+  ) where T: FixedWidthInteger, C: BinaryInteger {
     let floor = value >> count
     let lost = value &- floor << count
     let exact = count <= 0 || lost == 0
@@ -95,6 +95,10 @@ final class IntegerUtilitiesShiftTests: XCTestCase {
         testRoundingShift(T.random(in: .min ... .max), count, rounding: rule)
       }
     }
+    
+    for count in Int8.min ... .max {
+      testRoundingShift(T.random(in: .min ... .max), count, rounding: rule)
+    }
   }
 
   func testRoundingShifts() {
@@ -190,4 +194,55 @@ final class IntegerUtilitiesShiftTests: XCTestCase {
     testStochasticAverage(DoubleWidth<Int64>.random(in: .min ... .max))
     testStochasticAverage(DoubleWidth<UInt64>.random(in: .min ... .max))
   }
+  
+  func testSaturatingShift<T, C>(
+    _ value: T, _ count: C, rounding rule: RoundingRule
+  ) where T: FixedWidthInteger, C: FixedWidthInteger {
+    let observed = value.shiftedWithSaturation(leftBy: count, rounding: rule)
+    var expected: T = 0
+    if count <= 0 {
+      expected = value.shifted(rightBy: -Int64(count), rounding: rule)
+    } else {
+      let multiplier: T = 1 << count
+      if multiplier <= 0 {
+        expected = value == 0 ? 0 :
+                   value  < 0 ? .min : .max
+      } else {
+        expected = value.multipliedWithSaturation(by: multiplier)
+      }
+    }
+    if observed != expected {
+      print("Error found in \(T.self).shiftedWithSaturation(leftBy: \(count), rounding: \(rule)).")
+      print("   Value: \(String(value, radix: 16))")
+      print("Expected: \(String(expected, radix: 16))")
+      print("Observed: \(String(observed, radix: 16))")
+      XCTFail()
+      return
+    }
+  }
+  
+  func testSaturatingShift<T: FixedWidthInteger>(
+    _ type: T.Type, rounding rule: RoundingRule
+  ) {
+    for count in Int8.min ... .max {
+      testSaturatingShift(0, count, rounding: rule)
+      for bits in 0 ..< T.bitWidth {
+        let msb: T.Magnitude = 1 << bits
+        let value = T(truncatingIfNeeded: msb) | .random(in: 0 ... T(msb-1))
+        testSaturatingShift(value, count, rounding: rule)
+        testSaturatingShift(0 &- value, count, rounding: rule)
+      }
+    }
+  }
+  
+  func testSaturatingShifts() {
+    testSaturatingShift(Int8.self, rounding: .toOdd)
+    testSaturatingShift(UInt8.self, rounding: .toOdd)
+    testSaturatingShift(Int.self, rounding: .toOdd)
+    testSaturatingShift(UInt.self, rounding: .toOdd)
+  }
+  
+  func testEdgeCaseForNegativeCount() {
+    XCTAssertEqual(1.shiftedWithSaturation(leftBy: Int.min), 0)
+  }
 }