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