arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/src/dft/ntt.rs
src/dft/ntt.rsRaw
 1use crate::modulus::montgomery::Montgomery;
 2use crate::modulus::prime::Prime;
 3
 4pub struct Table<'a, O>{
 5    prime:&'a Prime<O>,
 6    pub psi_forward_rev:Vec<Montgomery<O>>,
 7    psi_backward_rev: Vec<Montgomery<O>>,
 8}
 9
10impl<'a> Table<'a, u64> {
11    pub fn new(prime: &'a mut Prime<u64>, nth_root: u64)->Self{
12
13        assert!(nth_root&(nth_root-1) == 0, "invalid argument: nth_root = {} is not a power of two", nth_root);
14
15        let psi: u64 = prime.primitive_nth_root(nth_root);
16
17        let psi_mont: Montgomery<u64> = prime.montgomery.prepare(psi);
18        let psi_inv_mont: Montgomery<u64> = prime.montgomery.pow(psi_mont, prime.phi-1);
19        
20        let mut psi_forward_rev: Vec<Montgomery<u64>> = vec![Montgomery(0); (nth_root >> 1) as usize];
21        let mut psi_backward_rev: Vec<Montgomery<u64>> = vec![Montgomery(0); (nth_root >> 1) as usize];
22
23        psi_forward_rev[0] = prime.montgomery.one();
24        psi_backward_rev[0] = prime.montgomery.one();
25
26        let log_nth_root_half: u32 = usize::BITS - ((nth_root>>1 as usize)-1).leading_zeros();
27
28        for i in 1..(nth_root>>1) as usize{
29
30            let i_rev_prev: usize = (i-1).reverse_bits() >> (usize::BITS - log_nth_root_half);
31            let i_rev_next: usize = i.reverse_bits() >> (usize::BITS - log_nth_root_half);
32
33            psi_forward_rev[i_rev_next] = prime.montgomery.mul_internal(psi_forward_rev[i_rev_prev], psi_mont);
34            psi_backward_rev[i_rev_next] = prime.montgomery.mul_internal(psi_backward_rev[i_rev_prev], psi_inv_mont);
35        }
36
37        Self{ 
38            prime: prime, 
39            psi_forward_rev: psi_forward_rev, 
40            psi_backward_rev: psi_backward_rev, 
41        }
42    }
43}