use square modulus

src/lib.rs

@@ -21,12 +21,25 @@ pub fn mod_exp(mut base: i64, mut exp: i64, p: i64) -> i64 {
 }
 
 fn mod_inv(a: i64, p: i64) -> i64 {
-    mod_exp(a, p - 2, p) // Using Fermat's Little Theorem
+    let sqrt_p = (p as f64).sqrt() as i64;
+    if sqrt_p * sqrt_p == p { 
+        // If p is a perfect square (p = q^2), use q^2 - q - 1
+        mod_exp(a, p - sqrt_p - 1, p)
+    } else {
+        // Otherwise, use standard Fermat’s theorem
+        mod_exp(a, p - 2, p)
+    }
 }
 
-// Compute n-th root of unity (omega = root^((p - 1) / n) % p)
-pub fn omega(root: i64, p: i64, n: usize) -> i64{
-    mod_exp(root, (p - 1) / n as i64, p)
+// Compute n-th root of unity (omega) for p, depending on whether p is a perfect square
+pub fn omega(root: i64, p: i64, n: usize) -> i64 {
+    // Check if p is a perfect square (p = q^2)
+    let sqrt_p = (p as f64).sqrt() as i64;
+    if sqrt_p * sqrt_p == p { 
+        mod_exp(root, (p - sqrt_p) / n as i64, p)  // order of mult. group is p - sqrt_p
+    } else {
+        mod_exp(root, (p - 1) / n as i64, p)  // order of mult. group is p - 1
+    }
 }
 
 // Forward transform using NTT, output bit-reversed 

src/test.rs

@@ -24,4 +24,27 @@ mod tests {
         // Ensure both methods produce the same result
         assert_eq!(c_std, c_fast, "The results of polymul and polymul_ntt do not match");
     }
+
+    #[test]
+    fn test_polymul_ntt_square_modulus() {
+        let p: i64 = 17; // Prime modulus
+        let root: i64 = 3; // Primitive root of unity
+        let n: usize = 8;  // Length of the NTT (must be a power of 2)
+        let omega = omega(root, p*p, n); // n-th root of unity
+
+        // Input polynomials (padded to length `n`)
+        let mut a = vec![1, 2, 3, 4];
+        let mut b = vec![5, 6, 7, 8];
+        a.resize(n, 0);
+        b.resize(n, 0);
+
+        // Perform the standard polynomial multiplication
+        let c_std = polymul(&a, &b, n as i64, p*p);
+
+        // Perform the NTT-based polynomial multiplication
+        let c_fast = polymul_ntt(&a, &b, n, p*p, omega);
+
+        // Ensure both methods produce the same result
+        assert_eq!(c_std, c_fast, "The results of polymul and polymul_ntt do not match");
+    }
 }