1#[cfg(feature = "parallel")]
2use rayon::prelude::*;
3#[cfg(feature = "parallel")]
4use rayon::slice::ParallelSliceMut;
5
6use crate::RingParams;
7use crate::modular::mod_pow;
8use crate::polys::Polynomial;
9
10#[derive(Clone)]
11pub struct NTTprecaculated {
12 pub twiddles: Box<[u128]>,
13 pub twiddles_inv: Box<[u128]>,
14}
15
16
17
18pub fn precalculate(params: &RingParams) -> NTTprecaculated {
19 let n = params.n;
20 let q = params.q as u128;
21 let psi = params.omega as u128;
22 let log_n = n.trailing_zeros();
23
24 let mut twiddles = vec![1u128; n];
25 for k in 0..n {
26 let br_k = k.reverse_bits() >> (64 - log_n);
27 let power = br_k as u128;
28 twiddles[k] = mod_pow(psi, power, q);
29 }
30
31
32 let psi_inv = mod_pow(psi, q - 2, q);
33 let mut twiddles_inv = vec![1u128; n];
34 for k in 0..n {
35 let br_k = k.reverse_bits() >> (64 - log_n);
36 let power = br_k as u128;
37 twiddles_inv[k] = mod_pow(psi_inv, power, q);
38 }
39
40 NTTprecaculated {
41 twiddles: twiddles.into_boxed_slice(),
42 twiddles_inv: twiddles_inv.into_boxed_slice(),
43 }
44}
45
46
47#[cfg(not(feature = "parallel"))]
48pub fn forward_ntt(params: &RingParams, polynomial: &Polynomial, ntt_tables: &NTTprecaculated) -> Polynomial {
49 forward_ntt_single( params, polynomial, ntt_tables)
50}
51
52#[cfg(not(feature = "parallel"))]
53pub fn inverse_ntt(params: &RingParams, polynomial: &Polynomial, ntt_tables: &NTTprecaculated) -> Polynomial {
54 inverse_ntt_single( params, polynomial, ntt_tables)
55}
56
57#[cfg(feature = "parallel")]
58pub fn forward_ntt(params: &RingParams, polynomial: &Polynomial, ntt_tables: &NTTprecaculated) -> Polynomial {
59 forward_ntt_multi( params, polynomial, ntt_tables)
60}
61
62#[cfg(feature = "parallel")]
63pub fn inverse_ntt(params: &RingParams, polynomial: &Polynomial, ntt_tables: &NTTprecaculated) -> Polynomial {
64 inverse_ntt_multi( params, polynomial, ntt_tables)
65}
66
67
68
69fn forward_ntt_single(
70 params: &RingParams,
71 polynomial: &Polynomial,
72 ntt_tables: &NTTprecaculated
73) -> Polynomial {
74 let mut coeffs: Vec<u64> = polynomial.coeffs.to_vec();
75
76
77 let n = params.n;
78 let q = params.q as u128;
79 let mut t: usize = n;
80 let mut m: usize = 1;
81
82 loop {
83 t /= 2;
84
85
86 for i in 0..m {
87 let j1 = 2 * i * t;
88 let j2 = j1 + t - 1;
89
90 for j in j1..=j2 {
91 let u = coeffs[j] as u128;
92 let w = ntt_tables.twiddles[m + i];
93 let v = ((coeffs[j + t]) as u128 * w) % q;
94
95 let mut x = u + v;
96 if x >= q { x -= q; }
97 coeffs[j] = x as u64;
98
99 let y = if u >= v { u - v } else { u + q - v };
100 coeffs[j + t] = y as u64;
101
102 }
103 }
104
105 m *= 2;
106
107 if m >= n { break; }
108 }
109
110 Polynomial {
111 coeffs: coeffs.into_iter().collect::<Vec<_>>().into_boxed_slice(),
112 }
113}
114
115fn inverse_ntt_single(
116 params: &RingParams,
117 polynomial: &Polynomial,
118 ntt_tables: &NTTprecaculated
119) -> Polynomial {
120 let mut coeffs: Vec<u64> = polynomial.coeffs.to_vec();
121
122 let n = params.n;
123 let q = params.q as u128;
124
125 let mut t: usize = 1;
126 let mut m: usize = n;
127
128 loop {
129 let h = m / 2;
130 let mut j1: usize = 0;
131
132 for i in 0..h {
133 let j2 = j1 + t - 1;
134
135 for j in j1..=j2 {
136 let u = coeffs[j] as u128;
137 let v = coeffs[j + t] as u128;
138
139 let sum = u + v;
140 let sum = if sum >= q { sum - q } else { sum };
141 coeffs[j] = sum as u64;
142
143
144 let w = ntt_tables.twiddles_inv[h + i];
145 coeffs[j + t] = (((u + q - v) % q * w) % q) as u64;
146 }
147
148 j1 += 2 * t;
149 }
150
151 t *= 2;
152 m /= 2;
153
154
155 if m <= 1 { break; }
156 }
157
158 let n_inv = mod_pow(n as u128, q - 2, q);
159 for c in coeffs.iter_mut() {
160 *c = (*c * n_inv as u64) % q as u64;
161 }
162
163 Polynomial {
164 coeffs: coeffs.into_iter().collect::<Vec<_>>().into_boxed_slice(),
165 }
166}
167
168
169#[cfg(feature = "parallel")]
170fn forward_ntt_multi(
171 params: &RingParams,
172 polynomial: &Polynomial,
173 ntt_tables: &NTTprecaculated
174) -> Polynomial {
175 let mut coeffs: Vec<u128> = polynomial.coeffs
176 .iter()
177 .map(|&c| c as u128)
178 .collect();
179
180 let n = params.n;
181 let q = params.q as u128;
182
183 let mut t: usize = n;
184 let mut m: usize = 1;
185
186 loop {
187 t = t / 2;
188
189
190
191 coeffs
192 .par_chunks_mut(2 * t)
193 .enumerate()
194 .for_each(|(i, chunk)| {
195 for j in 0..t {
196 let u = chunk[j];
197 let w = ntt_tables.twiddles[m + i];
198 let v = ((chunk[j + t]) * w) % params.q as u128;
199 let sum = u + v;
200
201 chunk[j] = if sum >= q { sum - q } else { sum };
202
203 chunk[j + t] = if u >= v { u - v } else { u + q - v };
204 }
205 });
206
207
208 m = m * 2;
209
210 if m >= n { break; }
211 }
212
213 Polynomial {
214 coeffs: coeffs.into_iter().map(|c| c as u64).collect::<Vec<_>>().into_boxed_slice(),
215 }
216}
217
218#[cfg(feature = "parallel")]
219fn inverse_ntt_multi(
220 params: &RingParams,
221 polynomial: &Polynomial,
222 ntt_tables: &NTTprecaculated
223) -> Polynomial {
224 let mut coeffs: Vec<u128> = polynomial.coeffs
225 .iter()
226 .map(|&c| c as u128)
227 .collect();
228
229 let n = params.n;
230 let q = params.q as u128;
231
232 let mut t: usize = 1;
233 let mut m: usize = n;
234
235 loop {
236 let h = m / 2;
237
238 coeffs
239 .par_chunks_mut(2 * t)
240 .enumerate()
241 .for_each(|(i, chunk)| {
242 let w = ntt_tables.twiddles_inv[h + i];
243
244 for j in 0..t {
245 let u = chunk[j];
246 let v = chunk[j + t];
247
248 let sum = u + v;
249 chunk[j] = if sum >= q { sum - q } else { sum };
250
251 chunk[j + t] = ((u + q - v) % q * w) % q;
252 }
253 });
254
255 t = t * 2;
256 m = m / 2;
257
258
259 if m <= 1 { break; }
260 }
261
262 let n_inv = mod_pow(n as u128, q - 2, q);
263 for c in coeffs.iter_mut() {
264 *c = (*c * n_inv) % q;
265 }
266
267 Polynomial {
268 coeffs: coeffs.into_iter().map(|c| c as u64).collect::<Vec<_>>().into_boxed_slice(),
269 }
270}