1//! Implementation of the NTT & iNTT, following the CT & GS algorighms, more
  2//! details in https://github.com/arnaucube/math/blob/master/notes_ntt.pdf .
  3use crate::zq::Zq;
  4
  5#[derive(Debug)]
  6pub struct NTT<const Q: u64, const N: usize> {}
  7
  8impl<const Q: u64, const N: usize> NTT<Q, N> {
  9    const N_INV: Zq<Q> = Zq(const_inv_mod::<Q>(N as u64));
 10    // since we work over Zq[X]/(X^N+1) (negacyclic), get the 2*N-th root of unity
 11    pub(crate) const ROOT_OF_UNITY: u64 = primitive_root_of_unity::<Q>(2 * N);
 12    pub(crate) const ROOTS_OF_UNITY: [Zq<Q>; N] = roots_of_unity(Self::ROOT_OF_UNITY);
 13    const ROOTS_OF_UNITY_INV: [Zq<Q>; N] = roots_of_unity_inv(Self::ROOTS_OF_UNITY);
 14}
 15
 16impl<const Q: u64, const N: usize> NTT<Q, N> {
 17    /// implements the Cooley-Tukey (CT) algorithm. Details at section 3.1 of
 18    /// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
 19    pub fn ntt(a: [Zq<Q>; N]) -> [Zq<Q>; N] {
 20        let mut t = N / 2;
 21        let mut m = 1;
 22        let mut r: [Zq<Q>; N] = a.clone();
 23        while m < N {
 24            let mut k = 0;
 25            for i in 0..m {
 26                let S: Zq<Q> = Self::ROOTS_OF_UNITY[m + i];
 27                for j in k..k + t {
 28                    let U: Zq<Q> = r[j];
 29                    let V: Zq<Q> = r[j + t] * S;
 30                    r[j] = U + V;
 31                    r[j + t] = U - V;
 32                }
 33                k = k + 2 * t;
 34            }
 35            t /= 2;
 36            m *= 2;
 37        }
 38        r
 39    }
 40
 41    /// implements the Gentleman-Sande (GS) algorithm. Details at section 3.2 of
 42    /// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
 43    pub fn intt(a: [Zq<Q>; N]) -> [Zq<Q>; N] {
 44        let mut t = 1;
 45        let mut m = N / 2;
 46        let mut r: [Zq<Q>; N] = a.clone();
 47        while m > 0 {
 48            let mut k = 0;
 49            for i in 0..m {
 50                let S: Zq<Q> = Self::ROOTS_OF_UNITY_INV[m + i];
 51                for j in k..k + t {
 52                    let U: Zq<Q> = r[j];
 53                    let V: Zq<Q> = r[j + t];
 54                    r[j] = U + V;
 55                    r[j + t] = (U - V) * S;
 56                }
 57                k += 2 * t;
 58            }
 59            t *= 2;
 60            m /= 2;
 61        }
 62        for i in 0..N {
 63            r[i] = r[i] * Self::N_INV;
 64        }
 65        r
 66    }
 67}
 68
 69/// computes a primitive N-th root of unity using the method described by Thomas
 70/// Pornin in https://crypto.stackexchange.com/a/63616
 71const fn primitive_root_of_unity<const Q: u64>(N: usize) -> u64 {
 72    assert!(N.is_power_of_two());
 73    assert!((Q - 1) % N as u64 == 0);
 74
 75    let n: u64 = N as u64;
 76    let mut k = 1;
 77    while k < Q {
 78        // alternatively could get a random k at each iteration, if so, add the following if:
 79        // `if k == 0 { continue; }`
 80        let w = const_exp_mod::<Q>(k, (Q - 1) / n);
 81        if const_exp_mod::<Q>(w, n / 2) != 1 {
 82            return w; // w is a primitive N-th root of unity
 83        }
 84        k += 1;
 85    }
 86    panic!("No primitive root of unity");
 87}
 88
 89const fn roots_of_unity<const Q: u64, const N: usize>(w: u64) -> [Zq<Q>; N] {
 90    let mut r: [Zq<Q>; N] = [Zq(0u64); N];
 91    let mut i = 0;
 92    let log_n = N.ilog2();
 93    while i < N {
 94        // (return the roots in bit-reverset order)
 95        let j = ((i as u64).reverse_bits() >> (64 - log_n)) as usize;
 96        r[i] = Zq(const_exp_mod::<Q>(w, j as u64));
 97        i += 1;
 98    }
 99    r
100}
101
102const fn roots_of_unity_inv<const Q: u64, const N: usize>(v: [Zq<Q>; N]) -> [Zq<Q>; N] {
103    // assumes that the inputted roots are already in bit-reverset order
104    let mut r: [Zq<Q>; N] = [Zq(0u64); N];
105    let mut i = 0;
106    while i < N {
107        r[i] = Zq(const_inv_mod::<Q>(v[i].0));
108        i += 1;
109    }
110    r
111}
112
113/// returns x^k mod Q
114const fn const_exp_mod<const Q: u64>(x: u64, k: u64) -> u64 {
115    let mut r = 1u64;
116    let mut x = x;
117    let mut k = k;
118    x = x % Q;
119    // exponentiation by square strategy
120    while k > 0 {
121        if k % 2 == 1 {
122            r = (r * x) % Q;
123        }
124        x = (x * x) % Q;
125        k /= 2;
126    }
127    r
128}
129
130/// returns x^-1 mod Q
131const fn const_inv_mod<const Q: u64>(x: u64) -> u64 {
132    // by Fermat's Little Theorem, x^-1 mod q \equiv  x^{q-2} mod q
133    const_exp_mod::<Q>(x, Q - 2)
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    use anyhow::Result;
141    use std::array;
142
143    #[test]
144    fn test_ntt() -> Result<()> {
145        const Q: u64 = 2u64.pow(16) + 1;
146        const N: usize = 4;
147
148        let a: [u64; N] = [1u64, 2, 3, 4];
149        let a: [Zq<Q>; N] = array::from_fn(|i| Zq::new(a[i]));
150
151        let a_ntt = NTT::<Q, N>::ntt(a);
152
153        let a_intt = NTT::<Q, N>::intt(a_ntt);
154
155        dbg!(&a);
156        dbg!(&a_ntt);
157        dbg!(&a_intt);
158        dbg!(NTT::<Q, N>::ROOT_OF_UNITY);
159        dbg!(NTT::<Q, N>::ROOTS_OF_UNITY);
160
161        assert_eq!(a, a_intt);
162        Ok(())
163    }
164
165    #[test]
166    fn test_ntt_loop() -> Result<()> {
167        const Q: u64 = 2u64.pow(16) + 1;
168        const N: usize = 512;
169
170        use rand::distributions::Distribution;
171        use rand::distributions::Uniform;
172        let mut rng = rand::thread_rng();
173        let dist = Uniform::new(0_f64, Q as f64);
174
175        for _ in 0..100 {
176            let a: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng)));
177            let a_ntt = NTT::<Q, N>::ntt(a);
178            let a_intt = NTT::<Q, N>::intt(a_ntt);
179            assert_eq!(a, a_intt);
180        }
181        Ok(())
182    }
183}