From 8b91d37d85e7d07f1ee79359069e3a0d407e0e2e Mon Sep 17 00:00:00 2001
From: Jackson Walters <jacksonwalters@gmail.com>
Date: Fri, 7 Feb 2025 14:57:32 -0500
Subject: [PATCH 1/2] compute inverse mod p^2

add case to compute inverses modulo p^2 to enable computation of NTT in this case
---
 src/lib.rs  |  9 ++++++++-
 src/test.rs | 23 +++++++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/src/lib.rs b/src/lib.rs
index 8aa0600..c69db7c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -21,7 +21,14 @@ 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)
diff --git a/src/test.rs b/src/test.rs
index 697c8cb..db21ff6 100644
--- a/src/test.rs
+++ b/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");
+    }
 }

From 3e2c743bacadcd89fe1df08345ea43bfb4a2da44 Mon Sep 17 00:00:00 2001
From: Jackson Walters <jacksonwalters@gmail.com>
Date: Fri, 7 Feb 2025 15:12:08 -0500
Subject: [PATCH 2/2] update omega function

use correct order of multiplicative group
---
 src/lib.rs | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index c69db7c..8cafe7f 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -31,9 +31,15 @@ fn mod_inv(a: i64, p: i64) -> i64 {
     }
 }
 
-// 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