1//! this file implements the non-efficient NTT, which uses multiplication by the
  2//! Vandermonde matrix.
  3use crate::zq::Zq;
  4
  5use anyhow::{anyhow, Result};
  6
  7#[derive(Debug)]
  8pub struct NTT<const Q: u64, const N: usize> {
  9    pub primitive: Zq<Q>,
 10    // nth_roots: Vec<Zq<Q>>,
 11    pub ntt: Vec<Vec<Zq<Q>>>,
 12    pub intt: Vec<Vec<Zq<Q>>>,
 13}
 14
 15impl<const Q: u64, const N: usize> NTT<Q, N> {
 16    pub fn new() -> Result<Self> {
 17        // TODO change n to be u64 and ensure that is n<Q
 18        // note: `n` here is not the `N` from `(X^N+1)`
 19        // TODO: in fact n will be N (trait/struct param)
 20
 21        // let primitive = Self::get_primitive_root_of_unity((2 * N) as u64)?;
 22        let primitive = Self::get_primitive_root_of_unity((2 * N) as u64)?;
 23        // let mut nth_roots = vec![Zq(0); N];
 24        // let mut w_i = Zq(1);
 25        // for i in 0..N {
 26        //     w_i = w_i * primitive;
 27        //     nth_roots[i] = w_i;
 28        // }
 29        let ntt: Vec<Vec<Zq<Q>>> = Self::vandermonde(primitive);
 30        let intt = Self::invert_vandermonde(&ntt);
 31        Ok(Self {
 32            primitive,
 33            // nth_roots,
 34            ntt,
 35            intt,
 36        })
 37    }
 38    pub fn vandermonde(primitive: Zq<Q>) -> Vec<Vec<Zq<Q>>> {
 39        let mut v: Vec<Vec<Zq<Q>>> = vec![];
 40        let n = (2 * N) as u64;
 41        // let n = N as u64;
 42        for i in 0..n {
 43            let mut row: Vec<Zq<Q>> = vec![];
 44            let primitive_i = primitive.exp(Zq(i));
 45            let mut primitive_ij = Zq(1);
 46            for _ in 0..n {
 47                row.push(primitive_ij);
 48                primitive_ij = primitive_ij * primitive_i;
 49            }
 50            v.push(row);
 51        }
 52        v
 53    }
 54    // specifically for the Vandermonde matrix
 55    pub fn invert_vandermonde(v: &Vec<Vec<Zq<Q>>>) -> Vec<Vec<Zq<Q>>> {
 56        let n = 2 * N;
 57        // let n = N;
 58        let mut inv: Vec<Vec<Zq<Q>>> = vec![];
 59        for i in 0..n {
 60            let w_i = v[i][1]; // = w_i^1=w^i^1 = w^i
 61            let w_i_inv = w_i.inv();
 62            let mut row: Vec<Zq<Q>> = vec![];
 63            for j in 0..n {
 64                row.push(w_i_inv.exp(Zq(j as u64)) / Zq(n as u64));
 65            }
 66            inv.push(row);
 67        }
 68        inv
 69    }
 70
 71    pub fn get_primitive_root_of_unity(n: u64) -> Result<Zq<Q>> {
 72        // using the method described by Thomas Pornin in
 73        // https://crypto.stackexchange.com/a/63616
 74
 75        // assert!((Q - 1) % N as u64 == 0);
 76        assert!((Q - 1) % n == 0);
 77
 78        // TODO maybe not using Zq and using u64 directly
 79        let n = Zq(n);
 80        for k in 0..Q {
 81            if k == 0 {
 82                continue;
 83            }
 84            let g = Zq(k);
 85            //  g = F.random_element()
 86            if g == Zq(0) {
 87                continue;
 88            }
 89            let w = g.exp((-Zq(1)) / n);
 90            if w.exp(n / Zq(2)) != Zq(1) {
 91                // g is the generator
 92                return Ok(w);
 93            }
 94        }
 95        Err(anyhow!("can not find the primitive root of unity"))
 96    }
 97}
 98
 99#[cfg(test)]
100mod tests {
101    use super::*;
102    use rand_distr::Uniform;
103
104    use crate::ring::matrix_vec_product;
105    use crate::ring::Rq;
106
107    #[test]
108    fn roots_of_unity() -> Result<()> {
109        const Q: u64 = 12289;
110        const N: usize = 512;
111        let _ntt = NTT::<Q, N>::new()?;
112        Ok(())
113    }
114
115    #[test]
116    fn vandermonde_ntt() -> Result<()> {
117        const Q: u64 = 41;
118        const N: usize = 4;
119        let primitive = NTT::<Q, N>::get_primitive_root_of_unity((2 * N) as u64)?;
120        let v = NTT::<Q, N>::vandermonde(primitive);
121
122        // naively compute the Vandermonde matrix, and assert that the one from the method matches
123        // the naively obtained one
124        let n2 = (2 * N) as u64;
125        let mut v2: Vec<Vec<Zq<Q>>> = vec![];
126        for i in 0..n2 {
127            let mut row: Vec<Zq<Q>> = vec![];
128            for j in 0..n2 {
129                row.push(primitive.exp(Zq(i * j)));
130            }
131            v2.push(row);
132        }
133        assert_eq!(v, v2);
134
135        let v_inv = NTT::<Q, N>::invert_vandermonde(&v);
136
137        let mut rng = rand::thread_rng();
138        let uniform_distr = Uniform::new(0_f64, Q as f64);
139        let a = Rq::<Q, N>::rand_f64(&mut rng, uniform_distr)?;
140        // let a = PR::<Q, N>::new_from_u64(vec![36, 21, 9, 19]);
141
142        // let a_padded_coeffs: [Zq<Q>; 2 * N] =
143        //     std::array::from_fn(|i| if i < N { a.coeffs[i] } else { Zq::zero() });
144        let mut a_padded = a.coeffs.to_vec();
145        a_padded.append(&mut vec![Zq(0); N]);
146        // let a_ntt = a_padded.mul_by_matrix(&v)?;
147        let a_ntt = matrix_vec_product(&v, &a_padded)?;
148        let a_intt: Vec<Zq<Q>> = matrix_vec_product(&v_inv, &a_ntt)?;
149        assert_eq!(a_intt, a_padded);
150        let a_intt_arr: [Zq<Q>; N] = std::array::from_fn(|i| a_intt[i]);
151        assert_eq!(Rq::new(a_intt_arr, None), a);
152
153        Ok(())
154    }
155
156    #[test]
157    fn vec_by_ntt() -> Result<()> {
158        const Q: u64 = 257;
159        const N: usize = 4;
160        // let primitive = NTT::<Q, N>::get_primitive_root_of_unity((2*N) as u64)?;
161        let ntt = NTT::<Q, N>::new()?;
162
163        let a: Vec<Zq<Q>> = vec![256, 256, 256, 256, 0, 0, 0, 0]
164            .iter()
165            .map(|&e| Zq::from_u64(e))
166            .collect();
167        let a_ntt = matrix_vec_product(&ntt.ntt, &a)?;
168        let a_intt = matrix_vec_product(&ntt.intt, &a_ntt)?;
169        assert_eq!(a_intt, a);
170
171        Ok(())
172    }
173
174    #[test]
175    fn bench_ntt() -> Result<()> {
176        // const Q: u64 = 12289;
177        // const N: usize = 512;
178        const Q: u64 = 257;
179        const N: usize = 4;
180        // let primitive = NTT::<Q, N>::get_primitive_root_of_unity((2*N) as u64)?;
181        let ntt = NTT::<Q, N>::new()?;
182
183        let rng = rand::thread_rng();
184        let a = Rq::<Q, { 2 * N }>::rand_f64(rng, Uniform::new(0_f64, (Q - 1) as f64))?;
185        let a = a.coeffs;
186        dbg!(&a);
187        let a_ntt = matrix_vec_product(&ntt.ntt, &a.to_vec())?;
188        dbg!(&a_ntt);
189        let a_intt = matrix_vec_product(&ntt.intt, &a_ntt)?;
190        dbg!(&a_intt);
191        assert_eq!(a_intt, a);
192        // TODO bench
193
194        Ok(())
195    }
196}