fhe-study/arith/src/naive_ntt.rs
arith/src/naive_ntt.rsRaw
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 /// returns the Vandermonde matrix for the given primitive root of unity.
39 /// Vandermonde matrix: https://en.wikipedia.org/wiki/Vandermonde_matrix
40 pub fn vandermonde(primitive: Zq<Q>) -> Vec<Vec<Zq<Q>>> {
41 let mut v: Vec<Vec<Zq<Q>>> = vec![];
42 let n = (2 * N) as u64;
43 // let n = N as u64;
44 for i in 0..n {
45 let mut row: Vec<Zq<Q>> = vec![];
46 let primitive_i = primitive.exp(Zq(i));
47 let mut primitive_ij = Zq(1);
48 for _ in 0..n {
49 row.push(primitive_ij);
50 primitive_ij = primitive_ij * primitive_i;
51 }
52 v.push(row);
53 }
54 v
55 }
56 // specifically for the Vandermonde matrix
57 /// returns the inverse Vandermonde matrix
58 pub fn invert_vandermonde(v: &Vec<Vec<Zq<Q>>>) -> Vec<Vec<Zq<Q>>> {
59 let n = 2 * N;
60 // let n = N;
61 let mut inv: Vec<Vec<Zq<Q>>> = vec![];
62 for i in 0..n {
63 let w_i = v[i][1]; // = w_i^1=w^i^1 = w^i
64 let w_i_inv = w_i.inv();
65 let mut row: Vec<Zq<Q>> = vec![];
66 for j in 0..n {
67 row.push(w_i_inv.exp(Zq(j as u64)) / Zq(n as u64));
68 }
69 inv.push(row);
70 }
71 inv
72 }
73
74 /// computes a primitive N-th root of unity using the method described by
75 /// Thomas Pornin in https://crypto.stackexchange.com/a/63616
76 pub fn get_primitive_root_of_unity(n: u64) -> Result<Zq<Q>> {
77 // using the method described by Thomas Pornin in
78 // https://crypto.stackexchange.com/a/63616
79
80 // assert!((Q - 1) % N as u64 == 0);
81 assert!((Q - 1) % n == 0);
82
83 // TODO maybe not using Zq and using u64 directly
84 let n = Zq(n);
85 for k in 0..Q {
86 if k == 0 {
87 continue;
88 }
89 let g = Zq(k);
90 // g = F.random_element()
91 if g == Zq(0) {
92 continue;
93 }
94 let w = g.exp((-Zq(1)) / n);
95 if w.exp(n / Zq(2)) != Zq(1) {
96 // g is the generator
97 return Ok(w);
98 }
99 }
100 Err(anyhow!("can not find the primitive root of unity"))
101 }
102}
103
104#[cfg(test)]
105mod tests {
106 use super::*;
107 use rand_distr::Uniform;
108
109 use crate::ring_nq::matrix_vec_product;
110 use crate::ring_nq::Rq;
111
112 #[test]
113 fn roots_of_unity() -> Result<()> {
114 const Q: u64 = 12289;
115 const N: usize = 512;
116 let _ntt = NTT::<Q, N>::new()?;
117 Ok(())
118 }
119
120 #[test]
121 fn vandermonde_ntt() -> Result<()> {
122 const Q: u64 = 41;
123 const N: usize = 4;
124 let primitive = NTT::<Q, N>::get_primitive_root_of_unity((2 * N) as u64)?;
125 let v = NTT::<Q, N>::vandermonde(primitive);
126
127 // naively compute the Vandermonde matrix, and assert that the one from the method matches
128 // the naively obtained one
129 let n2 = (2 * N) as u64;
130 let mut v2: Vec<Vec<Zq<Q>>> = vec![];
131 for i in 0..n2 {
132 let mut row: Vec<Zq<Q>> = vec![];
133 for j in 0..n2 {
134 row.push(primitive.exp(Zq(i * j)));
135 }
136 v2.push(row);
137 }
138 assert_eq!(v, v2);
139
140 let v_inv = NTT::<Q, N>::invert_vandermonde(&v);
141
142 let mut rng = rand::rng();
143 let uniform_distr = Uniform::new(0_f64, Q as f64)?;
144 let a = Rq::<Q, N>::rand_f64(&mut rng, uniform_distr)?;
145 // let a = PR::<Q, N>::new_from_u64(vec![36, 21, 9, 19]);
146
147 // let a_padded_coeffs: [Zq<Q>; 2 * N] =
148 // std::array::from_fn(|i| if i < N { a.coeffs[i] } else { Zq::zero() });
149 let mut a_padded = a.coeffs.to_vec();
150 a_padded.append(&mut vec![Zq(0); N]);
151 // let a_ntt = a_padded.mul_by_matrix(&v)?;
152 let a_ntt = matrix_vec_product(&v, &a_padded)?;
153 let a_intt: Vec<Zq<Q>> = matrix_vec_product(&v_inv, &a_ntt)?;
154 assert_eq!(a_intt, a_padded);
155 let a_intt_arr: [Zq<Q>; N] = std::array::from_fn(|i| a_intt[i]);
156 assert_eq!(Rq::new(a_intt_arr, None), a);
157
158 Ok(())
159 }
160
161 #[test]
162 fn vec_by_ntt() -> Result<()> {
163 const Q: u64 = 257;
164 const N: usize = 4;
165 // let primitive = NTT::<Q, N>::get_primitive_root_of_unity((2*N) as u64)?;
166 let ntt = NTT::<Q, N>::new()?;
167
168 let a: Vec<Zq<Q>> = vec![256, 256, 256, 256, 0, 0, 0, 0]
169 .iter()
170 .map(|&e| Zq::from_u64(e))
171 .collect();
172 let a_ntt = matrix_vec_product(&ntt.ntt, &a)?;
173 let a_intt = matrix_vec_product(&ntt.intt, &a_ntt)?;
174 assert_eq!(a_intt, a);
175
176 Ok(())
177 }
178
179 #[test]
180 fn bench_ntt() -> Result<()> {
181 // const Q: u64 = 12289;
182 // const N: usize = 512;
183 const Q: u64 = 257;
184 const N: usize = 4;
185 // let primitive = NTT::<Q, N>::get_primitive_root_of_unity((2*N) as u64)?;
186 let ntt = NTT::<Q, N>::new()?;
187
188 let rng = rand::rng();
189 let a = Rq::<Q, { 2 * N }>::rand_f64(rng, Uniform::new(0_f64, (Q - 1) as f64))?;
190 let a = a.coeffs;
191 dbg!(&a);
192 let a_ntt = matrix_vec_product(&ntt.ntt, &a.to_vec())?;
193 dbg!(&a_ntt);
194 let a_intt = matrix_vec_product(&ntt.intt, &a_ntt)?;
195 dbg!(&a_intt);
196 assert_eq!(a_intt, a);
197 // TODO bench
198
199 Ok(())
200 }
201}