From bbc236f7c9afc7a156e14a0ca3459c5245cd587b Mon Sep 17 00:00:00 2001
From: Jackson Walters <jacksonwalters@gmail.com>
Date: Tue, 11 Feb 2025 14:51:48 -0500
Subject: [PATCH] use totient function for mult. group size

it is cleaner to use the Euler totient function to compute the size of the group. this works for any modulus, and so we can now perform the NTT for arbitrary rings Z/NZ as long as n divides the group size.
---
 Cargo.toml |  3 ++-
 src/lib.rs | 24 ++++++++----------------
 2 files changed, 10 insertions(+), 17 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 255fd2b..f821f9b 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,4 +9,5 @@ homepage = "https://github.com/lattice-based-cryptography/ntt"
 repository = "https://github.com/lattice-based-cryptography/ntt"
 
 [dependencies]
-polynomial-ring = "0.5.0"
\ No newline at end of file
+polynomial-ring = "0.5.0"
+reikna = "0.12.3"
\ No newline at end of file
diff --git a/src/lib.rs b/src/lib.rs
index 8cafe7f..38f5b39 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,3 +1,5 @@
+use reikna::totient::totient;
+
 // Modular arithmetic functions using i64
 fn mod_add(a: i64, b: i64, p: i64) -> i64 {
     (a + b) % p
@@ -20,26 +22,16 @@ pub fn mod_exp(mut base: i64, mut exp: i64, p: i64) -> i64 {
     result
 }
 
+//compute the modular inverse of a modulo p using Fermat's little theorem, p not necessarily prime
 fn mod_inv(a: i64, p: i64) -> i64 {
-    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)
-    }
+    mod_exp(a, totient(p as u64) as i64 - 1, p) // order of mult. group is Euler's totient function
 }
 
-// Compute n-th root of unity (omega) for p, depending on whether p is a perfect square
+// Compute n-th root of unity (omega) for p not necessarily prime
 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
-    }
+    let grp_size = totient(p as u64) as i64;
+    assert!(grp_size % n as i64 == 0, "{} does not divide {}", n, grp_size);
+    mod_exp(root, grp_size / n as i64, p) // order of mult. group is Euler's totient function
 }
 
 // Forward transform using NTT, output bit-reversed