use totient function for mult. group size

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"
+polynomial-ring = "0.5.0"
+reikna = "0.12.3"

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