poulpy/math/src/ring/impl_u64/ring.rs
math/src/ring/impl_u64/ring.rsRaw
1use crate::dft::ntt::Table;
2use crate::modulus::barrett::Barrett;
3use crate::modulus::montgomery::Montgomery;
4use crate::modulus::prime::Prime;
5use crate::modulus::VectorOperations;
6use crate::modulus::{BARRETT, REDUCEMOD};
7use crate::poly::Poly;
8use crate::ring::Ring;
9use crate::CHUNK;
10use num_bigint::BigInt;
11use num_traits::ToPrimitive;
12
13impl Ring<u64> {
14 pub fn new(n: usize, q_base: u64, q_power: usize) -> Self {
15 let prime: Prime<u64> = Prime::<u64>::new(q_base, q_power);
16 Self {
17 n: n,
18 modulus: prime.clone(),
19 cyclotomic_order: n << 1,
20 dft: Box::new(Table::<u64>::new(prime, n << 1)),
21 }
22 }
23
24 pub fn from_bigint(&self, coeffs: &[BigInt], step: usize, a: &mut Poly<u64>) {
25 assert!(
26 step <= a.n(),
27 "invalid step: step={} > a.n()={}",
28 step,
29 a.n()
30 );
31 assert!(
32 coeffs.len() <= a.n() / step,
33 "invalid coeffs: coeffs.len()={} > a.n()/step={}",
34 coeffs.len(),
35 a.n() / step
36 );
37 let q_big: BigInt = BigInt::from(self.modulus.q);
38 a.0.iter_mut()
39 .step_by(step)
40 .enumerate()
41 .for_each(|(i, v)| *v = (&coeffs[i] % &q_big).to_u64().unwrap());
42 }
43}
44
45impl Ring<u64> {
46 pub fn ntt_inplace<const LAZY: bool>(&self, poly: &mut Poly<u64>) {
47 match LAZY {
48 true => self.dft.forward_inplace_lazy(&mut poly.0),
49 false => self.dft.forward_inplace(&mut poly.0),
50 }
51 }
52
53 pub fn intt_inplace<const LAZY: bool>(&self, poly: &mut Poly<u64>) {
54 match LAZY {
55 true => self.dft.backward_inplace_lazy(&mut poly.0),
56 false => self.dft.backward_inplace(&mut poly.0),
57 }
58 }
59
60 pub fn ntt<const LAZY: bool>(&self, poly_in: &Poly<u64>, poly_out: &mut Poly<u64>) {
61 poly_out.0.copy_from_slice(&poly_in.0);
62 match LAZY {
63 true => self.dft.forward_inplace_lazy(&mut poly_out.0),
64 false => self.dft.forward_inplace(&mut poly_out.0),
65 }
66 }
67
68 pub fn intt<const LAZY: bool>(&self, poly_in: &Poly<u64>, poly_out: &mut Poly<u64>) {
69 poly_out.0.copy_from_slice(&poly_in.0);
70 match LAZY {
71 true => self.dft.backward_inplace_lazy(&mut poly_out.0),
72 false => self.dft.backward_inplace(&mut poly_out.0),
73 }
74 }
75}
76
77impl Ring<u64> {
78 #[inline(always)]
79 pub fn a_add_b_into_b<const REDUCE: REDUCEMOD>(&self, a: &Poly<u64>, b: &mut Poly<u64>) {
80 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
81 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
82 self.modulus
83 .va_add_vb_into_vb::<CHUNK, REDUCE>(&a.0, &mut b.0);
84 }
85
86 #[inline(always)]
87 pub fn a_add_b_into_c<const REDUCE: REDUCEMOD>(
88 &self,
89 a: &Poly<u64>,
90 b: &Poly<u64>,
91 c: &mut Poly<u64>,
92 ) {
93 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
94 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
95 debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
96 self.modulus
97 .va_add_vb_into_vc::<CHUNK, REDUCE>(&a.0, &b.0, &mut c.0);
98 }
99
100 #[inline(always)]
101 pub fn a_add_b_scalar_into_a<const REDUCE: REDUCEMOD>(&self, b: &u64, a: &mut Poly<u64>) {
102 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
103 self.modulus.va_add_sb_into_va::<CHUNK, REDUCE>(b, &mut a.0);
104 }
105
106 #[inline(always)]
107 pub fn a_add_b_scalar_into_c<const REDUCE: REDUCEMOD>(
108 &self,
109 a: &Poly<u64>,
110 b: &u64,
111 c: &mut Poly<u64>,
112 ) {
113 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
114 debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
115 self.modulus
116 .va_add_sb_into_vc::<CHUNK, REDUCE>(&a.0, b, &mut c.0);
117 }
118
119 #[inline(always)]
120 pub fn a_add_scalar_b_mul_c_scalar_barrett_into_a<const REDUCE: REDUCEMOD>(
121 &self,
122 b: &u64,
123 c: &Barrett<u64>,
124 a: &mut Poly<u64>,
125 ) {
126 debug_assert!(a.n() == self.n(), "b.n()={} != n={}", a.n(), self.n());
127 self.modulus
128 .va_add_sb_mul_sc_barrett_into_va::<CHUNK, REDUCE>(b, c, &mut a.0);
129 }
130
131 #[inline(always)]
132 pub fn add_scalar_then_mul_scalar_barrett<const REDUCE: REDUCEMOD>(
133 &self,
134 a: &Poly<u64>,
135 b: &u64,
136 c: &Barrett<u64>,
137 d: &mut Poly<u64>,
138 ) {
139 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
140 debug_assert!(d.n() == self.n(), "c.n()={} != n={}", d.n(), self.n());
141 self.modulus
142 .va_add_sb_mul_sc_barrett_into_vd::<CHUNK, REDUCE>(&a.0, b, c, &mut d.0);
143 }
144
145 #[inline(always)]
146 pub fn a_sub_b_into_b<const BRANGE: u8, const REDUCE: REDUCEMOD>(
147 &self,
148 a: &Poly<u64>,
149 b: &mut Poly<u64>,
150 ) {
151 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
152 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
153 self.modulus
154 .va_sub_vb_into_vb::<CHUNK, BRANGE, REDUCE>(&a.0, &mut b.0);
155 }
156
157 #[inline(always)]
158 pub fn a_sub_b_into_a<const BRANGE: u8, const REDUCE: REDUCEMOD>(
159 &self,
160 b: &Poly<u64>,
161 a: &mut Poly<u64>,
162 ) {
163 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
164 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
165 self.modulus
166 .va_sub_vb_into_va::<CHUNK, BRANGE, REDUCE>(&b.0, &mut a.0);
167 }
168
169 #[inline(always)]
170 pub fn a_sub_b_into_c<const BRANGE: u8, const REDUCE: REDUCEMOD>(
171 &self,
172 a: &Poly<u64>,
173 b: &Poly<u64>,
174 c: &mut Poly<u64>,
175 ) {
176 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
177 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
178 debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
179 self.modulus
180 .va_sub_vb_into_vc::<CHUNK, BRANGE, REDUCE>(&a.0, &b.0, &mut c.0);
181 }
182
183 #[inline(always)]
184 pub fn a_neg_into_b<const ARANGE: u8, const REDUCE: REDUCEMOD>(
185 &self,
186 a: &Poly<u64>,
187 b: &mut Poly<u64>,
188 ) {
189 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
190 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
191 self.modulus
192 .va_neg_into_vb::<CHUNK, ARANGE, REDUCE>(&a.0, &mut b.0);
193 }
194
195 #[inline(always)]
196 pub fn a_neg_into_a<const ARANGE: u8, const REDUCE: REDUCEMOD>(&self, a: &mut Poly<u64>) {
197 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
198 self.modulus
199 .va_neg_into_va::<CHUNK, ARANGE, REDUCE>(&mut a.0);
200 }
201
202 #[inline(always)]
203 pub fn a_prepare_montgomery_into_a<const REDUCE: REDUCEMOD>(
204 &self,
205 a: &mut Poly<Montgomery<u64>>,
206 ) {
207 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
208 self.modulus
209 .va_prepare_montgomery_into_va::<CHUNK, REDUCE>(&mut a.0);
210 }
211
212 #[inline(always)]
213 pub fn a_mul_b_montgomery_into_c<const REDUCE: REDUCEMOD>(
214 &self,
215 a: &Poly<Montgomery<u64>>,
216 b: &Poly<u64>,
217 c: &mut Poly<u64>,
218 ) {
219 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
220 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
221 debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
222 self.modulus
223 .va_mul_vb_montgomery_into_vc::<CHUNK, REDUCE>(&a.0, &b.0, &mut c.0);
224 }
225
226 #[inline(always)]
227 pub fn a_mul_b_montgomery_add_c_into_c<const REDUCE1: REDUCEMOD, const REDUCE2: REDUCEMOD>(
228 &self,
229 a: &Poly<Montgomery<u64>>,
230 b: &Poly<u64>,
231 c: &mut Poly<u64>,
232 ) {
233 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
234 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
235 debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
236 self.modulus
237 .va_mul_vb_montgomery_add_vc_into_vc::<CHUNK, REDUCE1, REDUCE2>(&a.0, &b.0, &mut c.0);
238 }
239
240 #[inline(always)]
241 pub fn a_mul_b_montgomery_into_a<const REDUCE: REDUCEMOD>(
242 &self,
243 b: &Poly<Montgomery<u64>>,
244 a: &mut Poly<u64>,
245 ) {
246 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
247 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
248 self.modulus
249 .va_mul_vb_montgomery_into_va::<CHUNK, REDUCE>(&b.0, &mut a.0);
250 }
251
252 #[inline(always)]
253 pub fn a_mul_b_scalar_into_c<const REDUCE: REDUCEMOD>(
254 &self,
255 a: &Poly<u64>,
256 b: &u64,
257 c: &mut Poly<u64>,
258 ) {
259 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
260 debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
261 self.modulus.va_mul_sb_barrett_into_vc::<CHUNK, REDUCE>(
262 &a.0,
263 &self.modulus.barrett.prepare(*b),
264 &mut c.0,
265 );
266 }
267
268 #[inline(always)]
269 pub fn a_mul_b_scalar_into_a<const REDUCE: REDUCEMOD>(&self, b: &u64, a: &mut Poly<u64>) {
270 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
271 self.modulus.va_mul_sb_barrett_into_va::<CHUNK, REDUCE>(
272 &self
273 .modulus
274 .barrett
275 .prepare(self.modulus.barrett.reduce::<BARRETT>(b)),
276 &mut a.0,
277 );
278 }
279
280 #[inline(always)]
281 pub fn a_mul_b_scalar_barrett_into_a<const REDUCE: REDUCEMOD>(
282 &self,
283 b: &Barrett<u64>,
284 a: &mut Poly<u64>,
285 ) {
286 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
287 self.modulus
288 .va_mul_sb_barrett_into_va::<CHUNK, REDUCE>(b, &mut a.0);
289 }
290
291 #[inline(always)]
292 pub fn a_mul_b_scalar_barrett_into_c<const REDUCE: REDUCEMOD>(
293 &self,
294 a: &Poly<u64>,
295 b: &Barrett<u64>,
296 c: &mut Poly<u64>,
297 ) {
298 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
299 self.modulus
300 .va_mul_sb_barrett_into_vc::<CHUNK, REDUCE>(&a.0, b, &mut c.0);
301 }
302
303 #[inline(always)]
304 pub fn a_sub_b_mul_c_scalar_barrett_into_d<const VBRANGE: u8, const REDUCE: REDUCEMOD>(
305 &self,
306 a: &Poly<u64>,
307 b: &Poly<u64>,
308 c: &Barrett<u64>,
309 d: &mut Poly<u64>,
310 ) {
311 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
312 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
313 debug_assert!(d.n() == self.n(), "d.n()={} != n={}", d.n(), self.n());
314 self.modulus
315 .va_sub_vb_mul_sc_barrett_into_vd::<CHUNK, VBRANGE, REDUCE>(&a.0, &b.0, c, &mut d.0);
316 }
317
318 #[inline(always)]
319 pub fn b_sub_a_mul_c_scalar_barrett_into_a<const BRANGE: u8, const REDUCE: REDUCEMOD>(
320 &self,
321 b: &Poly<u64>,
322 c: &Barrett<u64>,
323 a: &mut Poly<u64>,
324 ) {
325 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
326 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
327 self.modulus
328 .va_sub_vb_mul_sc_barrett_into_vb::<CHUNK, BRANGE, REDUCE>(&b.0, c, &mut a.0);
329 }
330
331 #[inline(always)]
332 pub fn a_sub_b_add_c_scalar_mul_d_scalar_barrett_into_e<
333 const BRANGE: u8,
334 const REDUCE: REDUCEMOD,
335 >(
336 &self,
337 a: &Poly<u64>,
338 b: &Poly<u64>,
339 c: &u64,
340 d: &Barrett<u64>,
341 e: &mut Poly<u64>,
342 ) {
343 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
344 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
345 debug_assert!(e.n() == self.n(), "e.n()={} != n={}", e.n(), self.n());
346 self.modulus
347 .vb_sub_va_add_sc_mul_sd_barrett_into_ve::<CHUNK, BRANGE, REDUCE>(
348 &a.0, &b.0, c, d, &mut e.0,
349 );
350 }
351
352 #[inline(always)]
353 pub fn b_sub_a_add_c_scalar_mul_d_scalar_barrett_into_a<
354 const BRANGE: u8,
355 const REDUCE: REDUCEMOD,
356 >(
357 &self,
358 b: &Poly<u64>,
359 c: &u64,
360 d: &Barrett<u64>,
361 a: &mut Poly<u64>,
362 ) {
363 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
364 debug_assert!(b.n() == self.n(), "b.n()={} != n={}", b.n(), self.n());
365 self.modulus
366 .vb_sub_va_add_sc_mul_sd_barrett_into_va::<CHUNK, BRANGE, REDUCE>(&b.0, c, d, &mut a.0);
367 }
368
369 pub fn a_rsh_scalar_b_mask_scalar_c_into_d(
370 &self,
371 a: &Poly<u64>,
372 b: &usize,
373 c: &u64,
374 d: &mut Poly<u64>,
375 ) {
376 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
377 debug_assert!(d.n() == self.n(), "d.n()={} != n={}", d.n(), self.n());
378 self.modulus
379 .va_rsh_sb_mask_sc_into_vd::<CHUNK>(&a.0, b, c, &mut d.0);
380 }
381
382 pub fn a_rsh_scalar_b_mask_scalar_c_add_d_into_d(
383 &self,
384 a: &Poly<u64>,
385 b: &usize,
386 c: &u64,
387 d: &mut Poly<u64>,
388 ) {
389 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
390 debug_assert!(d.n() == self.n(), "d.n()={} != n={}", d.n(), self.n());
391 self.modulus
392 .va_rsh_sb_mask_sc_add_vd_into_vd::<CHUNK>(&a.0, b, c, &mut d.0);
393 }
394
395 pub fn a_ith_digit_unsigned_base_scalar_b_into_c(
396 &self,
397 i: usize,
398 a: &Poly<u64>,
399 b: &usize,
400 c: &mut Poly<u64>,
401 ) {
402 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
403 debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
404 self.modulus
405 .va_ith_digit_unsigned_base_sb_into_vc::<CHUNK>(i, &a.0, b, &mut c.0);
406 }
407
408 pub fn a_ith_digit_signed_base_scalar_b_into_c<const BALANCED: bool>(
409 &self,
410 i: usize,
411 a: &Poly<u64>,
412 b: &usize,
413 carry: &mut Poly<u64>,
414 c: &mut Poly<u64>,
415 ) {
416 debug_assert!(a.n() == self.n(), "a.n()={} != n={}", a.n(), self.n());
417 debug_assert!(
418 carry.n() == self.n(),
419 "carry.n()={} != n={}",
420 carry.n(),
421 self.n()
422 );
423 debug_assert!(c.n() == self.n(), "c.n()={} != n={}", c.n(), self.n());
424 self.modulus
425 .va_ith_digit_signed_base_sb_into_vc::<CHUNK, BALANCED>(
426 i,
427 &a.0,
428 b,
429 &mut carry.0,
430 &mut c.0,
431 );
432 }
433}