Fix debug tests

Sources/HomomorphicEncryption/PolyRq/PolyRq+Ntt.swift

@@ -241,11 +241,11 @@ extension PolyContext {
         let rootOfUnityPowers = nttContext.rootOfUnityPowers
         // The forward butterfly transforms `x,y` in
         // `[0, k * modulus)` -> `[0, (k + 2) * modulus)`.
-        // We delay modular reduction until overflowing T, i.e.
-        // `(kMax + 2) * modulus > T.max`, so `kMax = floor(T.max / modulus) - 2
+        // We delay modular reduction until overflowing the maximum input to `T.subtractIfExceeds`, i.e.,
+        // we find the largest `kMax` such that `(kMax + 2) * modulus <= (Self.max >> 1) + modulus`.
+        // So, we have `kMax =  T.max / (2 * modulus) - 1`
         var lazyReductionCounter = -1 // k
-        // kMax
-        let maxLazyReductionCounter = modulusReduceFactor.singleWordModulus.factor.low &- 2
+        let maxLazyReductionCounter = T.max / (2 * modulus) - 1 // kMax
 
         func applyFinalStageOp(m: Int, op: (_ x: inout T, _ y: inout T) -> Void) {
             for i in 0..<m {
@@ -292,14 +292,17 @@ extension PolyContext {
             let timeToReduce = lazyReductionCounter > maxLazyReductionCounter
             if timeToReduce {
                 if t == 1 {
-                    lazyReductionCounter &-= 2
+                    // if lazyReductionCounter == 3, `subtractIfExceeds(twiceModulus)`
+                    // only ensures `x in [0, 2 * p - 1]`
+                    lazyReductionCounter = max(lazyReductionCounter - 2, 2)
                 } else {
                     lazyReductionCounter = 1
                 }
             }
             switch (t, timeToReduce) {
             case (1, true):
                 applyFinalStageOp(m: m) { x, _ in
+                    // Ensure butterfly doesn't overflow
                     x = x.subtractIfExceeds(twiceModulus)
                 }
             case (1, false):

Sources/HomomorphicEncryption/Scalar.swift

@@ -169,11 +169,12 @@ extension ScalarType {
     ///
     /// Computes a conditional subtraction, `if self >= modulus ? self - modulus : self`, which can be used for modular
     /// reduction of `self` from range `[0, 2 * modulus - 1]` to `[0, modulus - 1]`. The computation is constant-time.
+    /// `self` must be less than or equal to `(Self.max >> 1) + modulus``
     /// - Parameter modulus: Modulus.
     /// - Returns: `self >= modulus ? self - modulus : self`.
     @inlinable
     public func subtractIfExceeds(_ modulus: Self) -> Self {
-        assert(self < 2 * modulus)
+        assert(self <= (Self.max &>> 1) + modulus) // difference mask fails otherwise
         let difference = self &- modulus
         let mask = Self(0) &- (difference >> (bitWidth - 1))
         return difference &+ (modulus & mask)

Tests/HomomorphicEncryptionTests/NttTests.swift

@@ -106,6 +106,55 @@ final class NttTests: XCTestCase {
             ])
     }
 
+    func testNtt16() throws {
+        // modulus near top of range
+        try runNttTest(
+            moduli: [UInt32(536_870_849)],
+            coeffData: [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
+            evalData: [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
+        try runNttTest(
+            moduli: [UInt32(536_870_849)],
+            coeffData: [[1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]],
+            evalData: [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])
+        try runNttTest(
+            moduli: [UInt32(536_870_849)],
+            coeffData: [[
+                477_051_601,
+                421_524_611,
+                456_257_859,
+                247_136_825,
+                128_775_020,
+                76_785_070,
+                49_764_016,
+                525_812_772,
+                325_605_371,
+                88_935_943,
+                255_470_762,
+                39_507_048,
+                404_978_219,
+                379_383_003,
+                244_420_585,
+                346_826_612,
+            ]], evalData: [[
+                230_846_094,
+                480_599_401,
+                157_364_576,
+                360_442_736,
+                531_052_463,
+                294_311_347,
+                432_899_854,
+                219_721_533,
+                286_807_067,
+                260_650_843,
+                362_842_688,
+                315_862_017,
+                493_042_020,
+                520_739_674,
+                167_758_416,
+                370_401_491,
+            ]])
+    }
+
     func testNtt32() throws {
         let modulus = UInt32(769)
 
@@ -175,8 +224,11 @@ final class NttTests: XCTestCase {
             return try xEval.inverseNtt()
         }
 
-        let moduli = [UInt64(576_460_752_303_436_801)]
         let degree = 128
+        let moduli = try UInt32.generatePrimes(
+            significantBitCounts: [30],
+            preferringSmall: false,
+            nttDegree: degree)
         let context = try PolyContext(degree: degree, moduli: moduli)
         let x = PolyRq<_, Coeff>.random(context: context)
         let y = PolyRq<_, Coeff>.random(context: context)

Tests/HomomorphicEncryptionTests/ScalarTests.swift

@@ -17,6 +17,20 @@ import TestUtilities
 import XCTest
 
 class ScalarTests: XCTestCase {
+    func testSubtractIfExceeds() {
+        do {
+            let modulus: UInt32 = (1 << 29) - 63
+            XCTAssertEqual(UInt32(2 * modulus + 1).subtractIfExceeds(modulus), modulus + 1)
+            XCTAssertEqual(UInt32(modulus - 1).subtractIfExceeds(modulus), modulus - 1)
+        }
+        do {
+            let modulus: UInt32 = (1 << 31) - 10
+            let max = (UInt32.max >> 1) + modulus
+            XCTAssertEqual(UInt32(max).subtractIfExceeds(modulus), max - modulus)
+            XCTAssertEqual(UInt32(modulus - 1).subtractIfExceeds(modulus), modulus - 1)
+        }
+    }
+
     func testAddMod() {
         XCTAssertEqual(UInt32(0).addMod(1, modulus: 3), 1)
         XCTAssertEqual(UInt32(1).addMod(2, modulus: 3), 0)