From 9f3c902d9254f4e77787c390094e83c98f95eb13 Mon Sep 17 00:00:00 2001
From: Jackson Walters <jacksonwalters@gmail.com>
Date: Sun, 16 Feb 2025 20:50:33 -0500
Subject: [PATCH] use large n

---
 src/lib.rs  |  2 --
 src/test.rs | 29 +++++++++++++++++------------
 2 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 9237863..29d6f1c 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -156,11 +156,9 @@ fn factorize(n: i64) -> HashMap<i64, u32> {
 
 /// Fast computation of a primitive root mod p^e
 pub fn primitive_root(p: i64, e: u32) -> i64 {
-    println!("primitive_root called");
     let g = primitive_root_mod_p(p);
     let mut g_lifted = g; // Lift it to p^e
     for _ in 1..e {
-        println!("g_lifted: {}", g_lifted);
         if mod_exp(g_lifted, p-1, p.pow(e)) == 1 {
             g_lifted += p.pow(e - 1);
         }
diff --git a/src/test.rs b/src/test.rs
index 1a6edfe..c02a031 100644
--- a/src/test.rs
+++ b/src/test.rs
@@ -26,19 +26,24 @@ mod tests {
 
     #[test]
     fn test_polymul_ntt_square_modulus() {
-        let moduli = [17*17, 12289*12289]; // Different moduli to test
-        let n: usize = 8;  // Length of the NTT (must be a power of 2)
-
-        for &modulus in &moduli {
-            let omega = omega(modulus, n); // n-th root of unity
-            let mut a = vec![1, 2, 3, 4];
-            let mut b = vec![5, 6, 7, 8];
-            a.resize(n, 0);
-            b.resize(n, 0);
-            let c_std = polymul(&a, &b, n as i64, modulus);
-            let c_fast = polymul_ntt(&a, &b, n, modulus, omega);
-            assert_eq!(c_std, c_fast, "The results of polymul and polymul_ntt do not match");
+        let cases = [
+            (17*17, 4),        // small square modulus
+            (12289*12289, 512) // large square modulus
+        ];
+        
+        for &(modulus, n) in &cases {
+            let omega = omega(modulus, 2*n); // n-th root of unity
+            let mut a: Vec<i64> = (0..n).map(|x| x as i64).collect();
+            let mut b: Vec<i64> = (0..n).map(|x| x as i64).collect();
+            a.resize(2*n, 0);
+            b.resize(2*n, 0);
+        
+            let c_std = polymul(&a, &b, 2*n as i64, modulus);
+            let c_fast = polymul_ntt(&a, &b, 2*n, modulus, omega);
+        
+            assert_eq!(c_std, c_fast, "The results of polymul and polymul_ntt do not match for modulus {} and n {}", modulus, n);
         }
+
     }
 
     #[test]