negacyclic ntt works

Author: mhchia

src/ntt.rs

@@ -1,20 +1,4 @@
-use core::panic;
-
-use super::zq::Zq;
-/*
-src/
-  poly.rs    ← Poly, Rq, RqNtt
-  ntt.rs     ← ntt(), intt() methods
-
-Rq <-> RqNTT with methods
-
-impl Rq<Q, D> {
-    fn ntt(self) -> RqNtt<Q, D> { ... }
-}
-impl RqNtt<Q, D> {
-    fn intt(self) -> Rq<Q, D> { ... }
-}
- */
+use super::{zq::Zq, poly::Poly};
 
 pub fn prime_factors(mut n: u64) -> Vec<u64> {
     let mut factors = Vec::new();
@@ -34,11 +18,13 @@ pub fn prime_factors(mut n: u64) -> Vec<u64> {
     factors
 }
 
+
+/// R_q = Z_q[X]/(X^d+1). X^d + 1 = 0 -> X^d = -1 \mod q
+/// -> X^{2d} = 1. Assume q is prime, Z_q^* is a cyclic group with order q-1
+/// i.e. \forall g \in Z_q^*, g^{q-1} = 1. Since g^{(q-1)/(2d)}^{2d} = 1,
+/// for g to exist 2d must divide (q-1).
 pub fn find_primitive_2d_root_of_unity<const Q: u64>(d: u64) -> Zq<Q> {
-    // R_q = Z_q[X]/(X^d+1). X^d + 1 = 0 -> X^d = -1 \mod q
-    // -> X^{2d} = 1. Assume q is prime, Z_q^* is a cyclic group with order q-1
-    // i.e. \forall g \in Z_q^*, g^{q-1} = 1. Since g^{(q-1)/(2d)}^{2d} = 1,
-    // for g to exist 2d must divide (q-1).
+
     let order = Q - 1;
     assert_eq!(
         order % (2 * d),
@@ -54,7 +40,6 @@ pub fn find_primitive_2d_root_of_unity<const Q: u64>(d: u64) -> Zq<Q> {
     let factors = prime_factors(order);
     for i in 2..order {
         let g = Zq::<Q>::new(i);
-        // o(g) =
         let is_generator = factors.iter().all(|&p| g.pow(order / p) != Zq::<Q>::one());
         if is_generator {
             // if g is a generator -> o(g) = q-1 -> o(g^{(q-1)/2d}) = 2d
@@ -64,17 +49,108 @@ pub fn find_primitive_2d_root_of_unity<const Q: u64>(d: u64) -> Zq<Q> {
     panic!("no multiplicative generator found for Q={Q}")
 }
 
+
+/// NTT: here we assume coeffs can be split *completely* for simplicity and efficiency
+/// in the split fields.
+/// This is implemented according to this great article https://electricdusk.com/ntt.html
+/// This requires {d} to be a power of two.
+pub fn ntt<const Q: u64>(coeffs: Vec<Zq<Q>>, psi: Zq<Q>, psi_power: u64) -> Vec<Zq<Q>> {
+    let d = coeffs.len();
+    assert!(d.is_power_of_two(), "d should be power of two to split completely: d={d}");
+    assert!((Q-1).is_multiple_of(2*d as u64));
+
+    // Terminal condition: when d = 1, it's the last split. Just returns
+    // the constant term.
+    if d == 1 {
+        return vec![coeffs[0]];
+    }
+
+    // E.g. d=256, root here is \psi^{128} since X^{256}+1 = (X^{128} - 1)(X^{128} + 1)
+    let root= psi.pow(psi_power);
+    // Here is the "butterfly" part
+    // E.g. we're at a \in Z_q[X] / (X^256+1) and we're gonna split to
+    // a_l \in Z_q[X]/(X^128 - \psi^128), a_r \in Z_q[X]/(X^128 + \psi^128).
+    // We just let replace all X^128=\psi^128  in a to become a_l,
+    //                         X^128=-\psi^128 in a to become a_r.
+    // Then,
+    //      a_l[0] = a[0] + psi^{128} * a[128]
+    //      a_r[0] = a[0] - psi^{128} * a[128]
+    // Since `a[0]` and `psi^{128} * a[128]` are reused for a_l and a_r, just different
+    // operator before the latter term.
+    // We can draw it as a butterfly.
+
+    let mut a_l: Vec<Zq<Q>> = Vec::new();
+    let mut a_r: Vec<Zq<Q>> = Vec::new();
+
+    for i in 0..(d/2) {
+        a_l.push(coeffs[i] + root * coeffs[i + d/2]);
+        a_r.push(coeffs[i] - root * coeffs[i + d/2]);
+    }
+
+    // Split the left/right poly all the way down and get the results.
+    let a_l_coeffs = ntt(a_l, psi, psi_power / 2);
+    let a_r_coeffs = ntt(a_r, psi, psi_power / 2 + (d/2) as u64);
+    a_l_coeffs.into_iter().chain(a_r_coeffs).collect()
+}
+
+pub fn intt<const Q: u64>(evals: Vec<Zq<Q>>) -> Vec<Zq<Q>> {
+    todo!()
+}
+
 #[cfg(test)]
 mod tests {
+
     use super::*;
 
     const Q: u64 = 17;
     const D: u64 = 4;
+    type F = Zq<Q>;
+
+    fn setup() -> Zq<Q> {
+        let psi = find_primitive_2d_root_of_unity::<Q>(D);
+        println!("psi={:?}", psi);
+        psi
+    }
 
     #[test]
     fn test_primitive_2d_root_of_unity() {
-        let omega = find_primitive_2d_root_of_unity::<Q>(D);
-        assert_eq!(omega.pow(2 * D), Zq::<Q>::one()); // w^{2d} = 1
-        assert_eq!(omega.pow(D), -Zq::<Q>::one()); // w^d = -1
+        let psi = setup();
+        assert_eq!(psi.pow(2 * D), F::one()); // w^{2d} = 1
+        assert_eq!(psi.pow(D), -F::one()); // w^d = -1
+    }
+
+    // Sage test vectors: q=17, d=4, negacyclic NTT (X^d+1)
+    // coeffs [16, 3, 0, 14] <-> evals [15, 0, 0, 15]
+    #[test]
+    fn test_ntt_forward() {
+        let psi = setup();
+        let d = 4;
+        let coeffs = vec![F::new(16), F::new(3), F::new(0), F::new(14)];
+        let expected_evals = vec![F::new(15), F::new(0), F::new(0), F::new(15)];
+
+        let odd_powers: Vec<_> = (0..d).map(|k| psi.pow(2*k as u64 + 1)).collect();
+        println!("roots: {:?}", odd_powers);
+
+        let a = Poly::new(coeffs.clone());
+        let evals: Vec<_> = odd_powers.iter().map(|w| a.clone().eval(w.value())).collect();
+        println!("evals: {:?}", evals);
+        assert_eq!(ntt::<Q>(coeffs, psi, d/2), expected_evals);
+    }
+
+
+    #[test]
+    fn test_intt_backward() {
+        let evals = vec![F::new(15), F::new(0), F::new(0), F::new(15)];
+        let expected_coeffs = vec![F::new(16), F::new(3), F::new(0), F::new(14)];
+        assert_eq!(intt::<Q>(evals), expected_coeffs);
+    }
+
+    #[test]
+    fn test_ntt_intt_roundtrip() {
+        type F = Zq<Q>;
+        let psi = setup();
+        let d = 4;
+        let coeffs = vec![F::new(16), F::new(3), F::new(0), F::new(14)];
+        // assert_eq!(intt::<Q>(ntt::<Q>(coeffs, psi, d)), coeffs);
     }
 }

src/poly.rs

@@ -24,6 +24,14 @@ impl<const Q: u64> Poly<Q> {
             coeffs: vec![Zq::one()],
         }
     }
+
+    pub fn eval(&self, x: u64) -> Zq<Q> {
+        let mut s = Zq::<Q>::zero();
+        for (i, &c) in self.coeffs.iter().enumerate() {
+            s = s + c * Zq::<Q>::new(x).pow(i as u64);
+        }
+        s
+    }
 }
 
 impl<const Q: u64> Add for Poly<Q> {
@@ -353,6 +361,42 @@ mod tests {
         assert_eq!(p(&[1, 2, 3]) * p(&[2]), p(&[2, 4, 6]));
     }
 
+    // ─── Poly::eval tests ───
+
+    #[test]
+    fn test_eval_constant() {
+        // f = 5, f(x) = 5 for all x
+        assert_eq!(p(&[5]).eval(0), F::new(5));
+        assert_eq!(p(&[5]).eval(3), F::new(5));
+    }
+
+    #[test]
+    fn test_eval_linear() {
+        // f = 3 + 5x
+        // f(0) = 3, f(1) = 8, f(2) = 13
+        assert_eq!(p(&[3, 5]).eval(0), F::new(3));
+        assert_eq!(p(&[3, 5]).eval(1), F::new(8));
+        assert_eq!(p(&[3, 5]).eval(2), F::new(13));
+    }
+
+    #[test]
+    fn test_eval_quadratic() {
+        // f = 3 + 5x + 2x^2  (mod 17)
+        // f(0) = 3
+        // f(1) = 3 + 5 + 2 = 10
+        // f(2) = 3 + 10 + 8 = 21 mod 17 = 4
+        // f(16) = f(-1 mod 17) = 3 - 5 + 2 = 0
+        assert_eq!(p(&[3, 5, 2]).eval(0), F::new(3));
+        assert_eq!(p(&[3, 5, 2]).eval(1), F::new(10));
+        assert_eq!(p(&[3, 5, 2]).eval(2), F::new(4));
+        assert_eq!(p(&[3, 5, 2]).eval(16), F::new(0));
+    }
+
+    #[test]
+    fn test_eval_zero_poly() {
+        assert_eq!(R::zero().eval(5), F::zero());
+    }
+
     // ─── Rq tests: R_q = Z_q[X]/(X^4 + 1), q=17, d=4 ───
 
     const D: usize = 4;