poulpy/base2k/src/vec_znx_ops.rs
base2k/src/vec_znx_ops.rsRaw
1use crate::ffi::module::MODULE;
2use crate::ffi::vec_znx;
3use crate::{Backend, Module, VecZnx, ZnxBase, ZnxBasics, ZnxInfos, ZnxLayout, switch_degree, znx_post_process_ternary_op};
4use std::cmp::min;
5pub trait VecZnxOps {
6 /// Allocates a new [VecZnx].
7 ///
8 /// # Arguments
9 ///
10 /// * `cols`: the number of polynomials.
11 /// * `size`: the number of size per polynomial (a.k.a small polynomials).
12 fn new_vec_znx(&self, cols: usize, size: usize) -> VecZnx;
13
14 fn new_vec_znx_from_bytes(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnx;
15 fn new_vec_znx_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnx;
16
17 /// Returns the minimum number of bytes necessary to allocate
18 /// a new [VecZnx] through [VecZnx::from_bytes].
19 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize;
20
21 fn vec_znx_normalize_tmp_bytes(&self, cols: usize) -> usize;
22
23 /// c <- a + b.
24 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
25
26 /// b <- b + a.
27 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx);
28
29 /// c <- a - b.
30 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
31
32 /// b <- a - b.
33 fn vec_znx_sub_ab_inplace(&self, b: &mut VecZnx, a: &VecZnx);
34
35 /// b <- b - a.
36 fn vec_znx_sub_ba_inplace(&self, b: &mut VecZnx, a: &VecZnx);
37
38 /// b <- -a.
39 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx);
40
41 /// b <- -b.
42 fn vec_znx_negate_inplace(&self, a: &mut VecZnx);
43
44 /// b <- a * X^k (mod X^{n} + 1)
45 fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
46
47 /// a <- a * X^k (mod X^{n} + 1)
48 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx);
49
50 /// b <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
51 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
52
53 /// a <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
54 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx);
55
56 /// Splits b into subrings and copies them them into a.
57 ///
58 /// # Panics
59 ///
60 /// This method requires that all [VecZnx] of b have the same ring degree
61 /// and that b.n() * b.len() <= a.n()
62 fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx);
63
64 /// Merges the subrings a into b.
65 ///
66 /// # Panics
67 ///
68 /// This method requires that all [VecZnx] of a have the same ring degree
69 /// and that a.n() * a.len() <= b.n()
70 fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>);
71}
72
73impl<B: Backend> VecZnxOps for Module<B> {
74 fn new_vec_znx(&self, cols: usize, size: usize) -> VecZnx {
75 VecZnx::new(self, cols, size)
76 }
77
78 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize {
79 VecZnx::bytes_of(self, cols, size)
80 }
81
82 fn new_vec_znx_from_bytes(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnx {
83 VecZnx::from_bytes(self, cols, size, bytes)
84 }
85
86 fn new_vec_znx_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnx {
87 VecZnx::from_bytes_borrow(self, cols, size, tmp_bytes)
88 }
89
90 fn vec_znx_normalize_tmp_bytes(&self, cols: usize) -> usize {
91 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(self.ptr) as usize * cols }
92 }
93
94 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
95 let op = ffi_ternary_op_factory(
96 self.ptr,
97 c.size(),
98 c.sl(),
99 a.size(),
100 a.sl(),
101 b.size(),
102 b.sl(),
103 vec_znx::vec_znx_add,
104 );
105 vec_znx_apply_binary_op::<B, false>(self, c, a, b, op);
106 }
107
108 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
109 unsafe {
110 let b_ptr: *mut VecZnx = b as *mut VecZnx;
111 Self::vec_znx_add(self, &mut *b_ptr, a, &*b_ptr);
112 }
113 }
114
115 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
116 let op = ffi_ternary_op_factory(
117 self.ptr,
118 c.size(),
119 c.sl(),
120 a.size(),
121 a.sl(),
122 b.size(),
123 b.sl(),
124 vec_znx::vec_znx_sub,
125 );
126 vec_znx_apply_binary_op::<B, true>(self, c, a, b, op);
127 }
128
129 fn vec_znx_sub_ab_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
130 unsafe {
131 let b_ptr: *mut VecZnx = b as *mut VecZnx;
132 Self::vec_znx_sub(self, &mut *b_ptr, a, &*b_ptr);
133 }
134 }
135
136 fn vec_znx_sub_ba_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
137 unsafe {
138 let b_ptr: *mut VecZnx = b as *mut VecZnx;
139 Self::vec_znx_sub(self, &mut *b_ptr, &*b_ptr, a);
140 }
141 }
142
143 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx) {
144 let op = ffi_binary_op_factory_type_0(
145 self.ptr,
146 b.size(),
147 b.sl(),
148 a.size(),
149 a.sl(),
150 vec_znx::vec_znx_negate,
151 );
152 vec_znx_apply_unary_op::<B>(self, b, a, op);
153 }
154
155 fn vec_znx_negate_inplace(&self, a: &mut VecZnx) {
156 unsafe {
157 let a_ptr: *mut VecZnx = a as *mut VecZnx;
158 Self::vec_znx_negate(self, &mut *a_ptr, &*a_ptr);
159 }
160 }
161
162 fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
163 let op = ffi_binary_op_factory_type_1(
164 self.ptr,
165 k,
166 b.size(),
167 b.sl(),
168 a.size(),
169 a.sl(),
170 vec_znx::vec_znx_rotate,
171 );
172 vec_znx_apply_unary_op::<B>(self, b, a, op);
173 }
174
175 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx) {
176 unsafe {
177 let a_ptr: *mut VecZnx = a as *mut VecZnx;
178 Self::vec_znx_rotate(self, k, &mut *a_ptr, &*a_ptr);
179 }
180 }
181
182 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each size.
183 ///
184 /// # Arguments
185 ///
186 /// * `a`: input.
187 /// * `b`: output.
188 /// * `k`: the power to which to map each coefficients.
189 /// * `a_size`: the number of a_size on which to apply the mapping.
190 ///
191 /// # Panics
192 ///
193 /// The method will panic if the argument `a` is greater than `a.size()`.
194 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
195 let op = ffi_binary_op_factory_type_1(
196 self.ptr,
197 k,
198 b.size(),
199 b.sl(),
200 a.size(),
201 a.sl(),
202 vec_znx::vec_znx_automorphism,
203 );
204 vec_znx_apply_unary_op::<B>(self, b, a, op);
205 }
206
207 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each size.
208 ///
209 /// # Arguments
210 ///
211 /// * `a`: input and output.
212 /// * `k`: the power to which to map each coefficients.
213 /// * `a_size`: the number of size on which to apply the mapping.
214 ///
215 /// # Panics
216 ///
217 /// The method will panic if the argument `size` is greater than `self.size()`.
218 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx) {
219 unsafe {
220 let a_ptr: *mut VecZnx = a as *mut VecZnx;
221 Self::vec_znx_automorphism(self, k, &mut *a_ptr, &*a_ptr);
222 }
223 }
224
225 fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx) {
226 let (n_in, n_out) = (a.n(), b[0].n());
227
228 debug_assert!(
229 n_out < n_in,
230 "invalid a: output ring degree should be smaller"
231 );
232 b[1..].iter().for_each(|bi| {
233 debug_assert_eq!(
234 bi.n(),
235 n_out,
236 "invalid input a: all VecZnx must have the same degree"
237 )
238 });
239
240 b.iter_mut().enumerate().for_each(|(i, bi)| {
241 if i == 0 {
242 switch_degree(bi, a);
243 self.vec_znx_rotate(-1, buf, a);
244 } else {
245 switch_degree(bi, buf);
246 self.vec_znx_rotate_inplace(-1, buf);
247 }
248 })
249 }
250
251 fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>) {
252 let (n_in, n_out) = (b.n(), a[0].n());
253
254 debug_assert!(
255 n_out < n_in,
256 "invalid a: output ring degree should be smaller"
257 );
258 a[1..].iter().for_each(|ai| {
259 debug_assert_eq!(
260 ai.n(),
261 n_out,
262 "invalid input a: all VecZnx must have the same degree"
263 )
264 });
265
266 a.iter().enumerate().for_each(|(_, ai)| {
267 switch_degree(b, ai);
268 self.vec_znx_rotate_inplace(-1, b);
269 });
270
271 self.vec_znx_rotate_inplace(a.len() as i64, b);
272 }
273}
274
275fn ffi_ternary_op_factory(
276 module_ptr: *const MODULE,
277 c_size: usize,
278 c_sl: usize,
279 a_size: usize,
280 a_sl: usize,
281 b_size: usize,
282 b_sl: usize,
283 op_fn: unsafe extern "C" fn(*const MODULE, *mut i64, u64, u64, *const i64, u64, u64, *const i64, u64, u64),
284) -> impl Fn(&mut [i64], &[i64], &[i64]) {
285 move |cv: &mut [i64], av: &[i64], bv: &[i64]| unsafe {
286 op_fn(
287 module_ptr,
288 cv.as_mut_ptr(),
289 c_size as u64,
290 c_sl as u64,
291 av.as_ptr(),
292 a_size as u64,
293 a_sl as u64,
294 bv.as_ptr(),
295 b_size as u64,
296 b_sl as u64,
297 )
298 }
299}
300
301fn ffi_binary_op_factory_type_0(
302 module_ptr: *const MODULE,
303 b_size: usize,
304 b_sl: usize,
305 a_size: usize,
306 a_sl: usize,
307 op_fn: unsafe extern "C" fn(*const MODULE, *mut i64, u64, u64, *const i64, u64, u64),
308) -> impl Fn(&mut [i64], &[i64]) {
309 move |bv: &mut [i64], av: &[i64]| unsafe {
310 op_fn(
311 module_ptr,
312 bv.as_mut_ptr(),
313 b_size as u64,
314 b_sl as u64,
315 av.as_ptr(),
316 a_size as u64,
317 a_sl as u64,
318 )
319 }
320}
321
322fn ffi_binary_op_factory_type_1(
323 module_ptr: *const MODULE,
324 k: i64,
325 b_size: usize,
326 b_sl: usize,
327 a_size: usize,
328 a_sl: usize,
329 op_fn: unsafe extern "C" fn(*const MODULE, i64, *mut i64, u64, u64, *const i64, u64, u64),
330) -> impl Fn(&mut [i64], &[i64]) {
331 move |bv: &mut [i64], av: &[i64]| unsafe {
332 op_fn(
333 module_ptr,
334 k,
335 bv.as_mut_ptr(),
336 b_size as u64,
337 b_sl as u64,
338 av.as_ptr(),
339 a_size as u64,
340 a_sl as u64,
341 )
342 }
343}
344
345#[inline(always)]
346pub fn vec_znx_apply_binary_op<B: Backend, const NEGATE: bool>(
347 module: &Module<B>,
348 c: &mut VecZnx,
349 a: &VecZnx,
350 b: &VecZnx,
351 op: impl Fn(&mut [i64], &[i64], &[i64]),
352) {
353 #[cfg(debug_assertions)]
354 {
355 assert_eq!(a.n(), module.n());
356 assert_eq!(b.n(), module.n());
357 assert_eq!(c.n(), module.n());
358 assert_ne!(a.as_ptr(), b.as_ptr());
359 }
360
361 let a_cols: usize = a.cols();
362 let b_cols: usize = b.cols();
363 let c_cols: usize = c.cols();
364
365 let min_ab_cols: usize = min(a_cols, b_cols);
366 let min_cols: usize = min(c_cols, min_ab_cols);
367
368 // Applies over shared cols between (a, b, c)
369 (0..min_cols).for_each(|i| op(c.at_poly_mut(i, 0), a.at_poly(i, 0), b.at_poly(i, 0)));
370 // Copies/Negates/Zeroes the remaining cols if op is not inplace.
371 if c.as_ptr() != a.as_ptr() && c.as_ptr() != b.as_ptr() {
372 znx_post_process_ternary_op::<VecZnx, NEGATE>(c, a, b);
373 }
374}
375
376#[inline(always)]
377pub fn vec_znx_apply_unary_op<B: Backend>(module: &Module<B>, b: &mut VecZnx, a: &VecZnx, op: impl Fn(&mut [i64], &[i64])) {
378 #[cfg(debug_assertions)]
379 {
380 assert_eq!(a.n(), module.n());
381 assert_eq!(b.n(), module.n());
382 }
383 let a_cols: usize = a.cols();
384 let b_cols: usize = b.cols();
385 let min_cols: usize = min(a_cols, b_cols);
386 // Applies over the shared cols between (a, b)
387 (0..min_cols).for_each(|i| op(b.at_poly_mut(i, 0), a.at_poly(i, 0)));
388 // Zeroes the remaining cols of b.
389 (min_cols..b_cols).for_each(|i| (0..b.size()).for_each(|j| b.zero_at(i, j)));
390}
391
392#[cfg(test)]
393mod tests {
394 use crate::{
395 Backend, FFT64, Module, Sampling, VecZnx, VecZnxOps, ZnxBasics, ZnxInfos, ZnxLayout, ffi::vec_znx,
396 znx_post_process_ternary_op,
397 };
398 use itertools::izip;
399 use sampling::source::Source;
400 use std::cmp::min;
401
402 #[test]
403 fn vec_znx_add() {
404 let n: usize = 8;
405 let module: Module<FFT64> = Module::<FFT64>::new(n);
406 let op = |cv: &mut [i64], av: &[i64], bv: &[i64]| {
407 izip!(cv.iter_mut(), bv.iter(), av.iter()).for_each(|(ci, bi, ai)| *ci = *bi + *ai);
408 };
409 test_binary_op::<false, _>(
410 &module,
411 &|c: &mut VecZnx, a: &VecZnx, b: &VecZnx| module.vec_znx_add(c, a, b),
412 op,
413 );
414 }
415
416 #[test]
417 fn vec_znx_add_inplace() {
418 let n: usize = 8;
419 let module: Module<FFT64> = Module::<FFT64>::new(n);
420 let op = |bv: &mut [i64], av: &[i64]| {
421 izip!(bv.iter_mut(), av.iter()).for_each(|(bi, ai)| *bi = *bi + *ai);
422 };
423 test_binary_op_inplace::<false, _>(
424 &module,
425 &|b: &mut VecZnx, a: &VecZnx| module.vec_znx_add_inplace(b, a),
426 op,
427 );
428 }
429
430 #[test]
431 fn vec_znx_sub() {
432 let n: usize = 8;
433 let module: Module<FFT64> = Module::<FFT64>::new(n);
434 let op = |cv: &mut [i64], av: &[i64], bv: &[i64]| {
435 izip!(cv.iter_mut(), bv.iter(), av.iter()).for_each(|(ci, bi, ai)| *ci = *bi - *ai);
436 };
437 test_binary_op::<true, _>(
438 &module,
439 &|c: &mut VecZnx, a: &VecZnx, b: &VecZnx| module.vec_znx_sub(c, a, b),
440 op,
441 );
442 }
443
444 #[test]
445 fn vec_znx_sub_ab_inplace() {
446 let n: usize = 8;
447 let module: Module<FFT64> = Module::<FFT64>::new(n);
448 let op = |bv: &mut [i64], av: &[i64]| {
449 izip!(bv.iter_mut(), av.iter()).for_each(|(bi, ai)| *bi = *ai - *bi);
450 };
451 test_binary_op_inplace::<true, _>(
452 &module,
453 &|b: &mut VecZnx, a: &VecZnx| module.vec_znx_sub_ab_inplace(b, a),
454 op,
455 );
456 }
457
458 #[test]
459 fn vec_znx_sub_ba_inplace() {
460 let n: usize = 8;
461 let module: Module<FFT64> = Module::<FFT64>::new(n);
462 let op = |bv: &mut [i64], av: &[i64]| {
463 izip!(bv.iter_mut(), av.iter()).for_each(|(bi, ai)| *bi = *bi - *ai);
464 };
465 test_binary_op_inplace::<false, _>(
466 &module,
467 &|b: &mut VecZnx, a: &VecZnx| module.vec_znx_sub_ba_inplace(b, a),
468 op,
469 );
470 }
471
472 #[test]
473 fn vec_znx_negate() {
474 let n: usize = 8;
475 let module: Module<FFT64> = Module::<FFT64>::new(n);
476 let op = |b: &mut [i64], a: &[i64]| {
477 izip!(b.iter_mut(), a.iter()).for_each(|(bi, ai)| *bi = -*ai);
478 };
479 test_unary_op(
480 &module,
481 |b: &mut VecZnx, a: &VecZnx| module.vec_znx_negate(b, a),
482 op,
483 )
484 }
485
486 #[test]
487 fn vec_znx_negate_inplace() {
488 let n: usize = 8;
489 let module: Module<FFT64> = Module::<FFT64>::new(n);
490 let op = |a: &mut [i64]| a.iter_mut().for_each(|xi| *xi = -*xi);
491 test_unary_op_inplace(
492 &module,
493 |a: &mut VecZnx| module.vec_znx_negate_inplace(a),
494 op,
495 )
496 }
497
498 #[test]
499 fn vec_znx_rotate() {
500 let n: usize = 8;
501 let module: Module<FFT64> = Module::<FFT64>::new(n);
502 let k: i64 = 53;
503 let op = |b: &mut [i64], a: &[i64]| {
504 assert_eq!(b.len(), a.len());
505 b.copy_from_slice(a);
506
507 let mut k_mod2n: i64 = k % (2 * n as i64);
508 if k_mod2n < 0 {
509 k_mod2n += 2 * n as i64;
510 }
511 let sign: i64 = (k_mod2n.abs() / (n as i64)) & 1;
512 let k_modn: i64 = k_mod2n % (n as i64);
513
514 b.rotate_right(k_modn as usize);
515 b[0..k_modn as usize].iter_mut().for_each(|x| *x = -*x);
516
517 if sign == 1 {
518 b.iter_mut().for_each(|x| *x = -*x);
519 }
520 };
521 test_unary_op(
522 &module,
523 |b: &mut VecZnx, a: &VecZnx| module.vec_znx_rotate(k, b, a),
524 op,
525 )
526 }
527
528 #[test]
529 fn vec_znx_rotate_inplace() {
530 let n: usize = 8;
531 let module: Module<FFT64> = Module::<FFT64>::new(n);
532 let k: i64 = 53;
533 let rot = |a: &mut [i64]| {
534 let mut k_mod2n: i64 = k % (2 * n as i64);
535 if k_mod2n < 0 {
536 k_mod2n += 2 * n as i64;
537 }
538 let sign: i64 = (k_mod2n.abs() / (n as i64)) & 1;
539 let k_modn: i64 = k_mod2n % (n as i64);
540
541 a.rotate_right(k_modn as usize);
542 a[0..k_modn as usize].iter_mut().for_each(|x| *x = -*x);
543
544 if sign == 1 {
545 a.iter_mut().for_each(|x| *x = -*x);
546 }
547 };
548 test_unary_op_inplace(
549 &module,
550 |a: &mut VecZnx| module.vec_znx_rotate_inplace(k, a),
551 rot,
552 )
553 }
554
555 #[test]
556 fn vec_znx_automorphism() {
557 let n: usize = 8;
558 let module: Module<FFT64> = Module::<FFT64>::new(n);
559 let k: i64 = -5;
560 let op = |b: &mut [i64], a: &[i64]| {
561 assert_eq!(b.len(), a.len());
562 unsafe {
563 vec_znx::vec_znx_automorphism(
564 module.ptr,
565 k,
566 b.as_mut_ptr(),
567 1u64,
568 n as u64,
569 a.as_ptr(),
570 1u64,
571 n as u64,
572 );
573 }
574 };
575 test_unary_op(
576 &module,
577 |b: &mut VecZnx, a: &VecZnx| module.vec_znx_automorphism(k, b, a),
578 op,
579 )
580 }
581
582 #[test]
583 fn vec_znx_automorphism_inplace() {
584 let n: usize = 8;
585 let module: Module<FFT64> = Module::<FFT64>::new(n);
586 let k: i64 = -5;
587 let op = |a: &mut [i64]| unsafe {
588 vec_znx::vec_znx_automorphism(
589 module.ptr,
590 k,
591 a.as_mut_ptr(),
592 1u64,
593 n as u64,
594 a.as_ptr(),
595 1u64,
596 n as u64,
597 );
598 };
599 test_unary_op_inplace(
600 &module,
601 |a: &mut VecZnx| module.vec_znx_automorphism_inplace(k, a),
602 op,
603 )
604 }
605
606 fn test_binary_op<const NEGATE: bool, B: Backend>(
607 module: &Module<B>,
608 func_have: impl Fn(&mut VecZnx, &VecZnx, &VecZnx),
609 func_want: impl Fn(&mut [i64], &[i64], &[i64]),
610 ) {
611 let a_size: usize = 3;
612 let b_size: usize = 4;
613 let c_size: usize = 5;
614 let mut source: Source = Source::new([0u8; 32]);
615
616 [1usize, 2, 3].iter().for_each(|a_cols| {
617 [1usize, 2, 3].iter().for_each(|b_cols| {
618 [1usize, 2, 3].iter().for_each(|c_cols| {
619 let min_ab_cols: usize = min(*a_cols, *b_cols);
620 let min_cols: usize = min(*c_cols, min_ab_cols);
621 let min_size: usize = min(c_size, min(a_size, b_size));
622
623 let mut a: VecZnx = module.new_vec_znx(*a_cols, a_size);
624 (0..*a_cols).for_each(|i| {
625 module.fill_uniform(3, &mut a, i, a_size, &mut source);
626 });
627
628 let mut b: VecZnx = module.new_vec_znx(*b_cols, b_size);
629 (0..*b_cols).for_each(|i| {
630 module.fill_uniform(3, &mut b, i, b_size, &mut source);
631 });
632
633 let mut c_have: VecZnx = module.new_vec_znx(*c_cols, c_size);
634 (0..c_have.cols()).for_each(|i| {
635 module.fill_uniform(3, &mut c_have, i, c_size, &mut source);
636 });
637
638 func_have(&mut c_have, &a, &b);
639
640 let mut c_want: VecZnx = module.new_vec_znx(*c_cols, c_size);
641
642 // Adds with the minimum matching columns
643 (0..min_cols).for_each(|i| {
644 // Adds with th eminimum matching size
645 (0..min_size).for_each(|j| {
646 func_want(c_want.at_poly_mut(i, j), b.at_poly(i, j), a.at_poly(i, j));
647 });
648
649 if a_size > b_size {
650 // Copies remaining size of lh if lh.size() > rh.size()
651 (min_size..a_size).for_each(|j| {
652 izip!(c_want.at_poly_mut(i, j).iter_mut(), a.at_poly(i, j).iter()).for_each(|(ci, ai)| *ci = *ai);
653 if NEGATE {
654 c_want.at_poly_mut(i, j).iter_mut().for_each(|x| *x = -*x);
655 }
656 });
657 } else {
658 // Copies the remaining size of rh if the are greater
659 (min_size..b_size).for_each(|j| {
660 izip!(c_want.at_poly_mut(i, j).iter_mut(), b.at_poly(i, j).iter()).for_each(|(ci, bi)| *ci = *bi);
661 if NEGATE {
662 c_want.at_poly_mut(i, j).iter_mut().for_each(|x| *x = -*x);
663 }
664 });
665 }
666 });
667
668 znx_post_process_ternary_op::<_, NEGATE>(&mut c_want, &a, &b);
669
670 assert_eq!(c_have.raw(), c_want.raw());
671 });
672 });
673 });
674 }
675
676 fn test_binary_op_inplace<const NEGATE: bool, B: Backend>(
677 module: &Module<B>,
678 func_have: impl Fn(&mut VecZnx, &VecZnx),
679 func_want: impl Fn(&mut [i64], &[i64]),
680 ) {
681 let a_size: usize = 3;
682 let b_size: usize = 5;
683 let mut source = Source::new([0u8; 32]);
684
685 [1usize, 2, 3].iter().for_each(|a_cols| {
686 [1usize, 2, 3].iter().for_each(|b_cols| {
687 let min_cols: usize = min(*b_cols, *a_cols);
688 let min_size: usize = min(b_size, a_size);
689
690 let mut a: VecZnx = module.new_vec_znx(*a_cols, a_size);
691 (0..*a_cols).for_each(|i| {
692 module.fill_uniform(3, &mut a, i, a_size, &mut source);
693 });
694
695 let mut b_have: VecZnx = module.new_vec_znx(*b_cols, b_size);
696 (0..*b_cols).for_each(|i| {
697 module.fill_uniform(3, &mut b_have, i, b_size, &mut source);
698 });
699
700 let mut b_want: VecZnx = module.new_vec_znx(*b_cols, b_size);
701 b_want.raw_mut().copy_from_slice(b_have.raw());
702
703 func_have(&mut b_have, &a);
704
705 // Applies with the minimum matching columns
706 (0..min_cols).for_each(|i| {
707 // Adds with th eminimum matching size
708 (0..min_size).for_each(|j| func_want(b_want.at_poly_mut(i, j), a.at_poly(i, j)));
709 if NEGATE {
710 (min_size..b_size).for_each(|j| {
711 b_want.at_poly_mut(i, j).iter_mut().for_each(|x| *x = -*x);
712 });
713 }
714 });
715
716 assert_eq!(b_have.raw(), b_want.raw());
717 });
718 });
719 }
720
721 fn test_unary_op<B: Backend>(
722 module: &Module<B>,
723 func_have: impl Fn(&mut VecZnx, &VecZnx),
724 func_want: impl Fn(&mut [i64], &[i64]),
725 ) {
726 let a_size: usize = 3;
727 let b_size: usize = 5;
728 let mut source = Source::new([0u8; 32]);
729
730 [1usize, 2, 3].iter().for_each(|a_cols| {
731 [1usize, 2, 3].iter().for_each(|b_cols| {
732 let min_cols: usize = min(*b_cols, *a_cols);
733 let min_size: usize = min(b_size, a_size);
734
735 let mut a: VecZnx = module.new_vec_znx(*a_cols, a_size);
736 (0..a.cols()).for_each(|i| {
737 module.fill_uniform(3, &mut a, i, a_size, &mut source);
738 });
739
740 let mut b_have: VecZnx = module.new_vec_znx(*b_cols, b_size);
741 (0..b_have.cols()).for_each(|i| {
742 module.fill_uniform(3, &mut b_have, i, b_size, &mut source);
743 });
744
745 let mut b_want: VecZnx = module.new_vec_znx(*b_cols, b_size);
746
747 func_have(&mut b_have, &a);
748
749 // Applies on the minimum matching columns
750 (0..min_cols).for_each(|i| {
751 // Applies on the minimum matching size
752 (0..min_size).for_each(|j| func_want(b_want.at_poly_mut(i, j), a.at_poly(i, j)));
753
754 // Zeroes the unmatching size
755 (min_size..b_size).for_each(|j| {
756 b_want.zero_at(i, j);
757 })
758 });
759
760 // Zeroes the unmatching columns
761 (min_cols..*b_cols).for_each(|i| {
762 (0..b_size).for_each(|j| {
763 b_want.zero_at(i, j);
764 })
765 });
766
767 assert_eq!(b_have.raw(), b_want.raw());
768 });
769 });
770 }
771
772 fn test_unary_op_inplace<B: Backend>(module: &Module<B>, func_have: impl Fn(&mut VecZnx), func_want: impl Fn(&mut [i64])) {
773 let a_size: usize = 3;
774 let mut source = Source::new([0u8; 32]);
775 [1usize, 2, 3].iter().for_each(|a_cols| {
776 let mut a_have: VecZnx = module.new_vec_znx(*a_cols, a_size);
777 (0..*a_cols).for_each(|i| {
778 module.fill_uniform(3, &mut a_have, i, a_size, &mut source);
779 });
780
781 let mut a_want: VecZnx = module.new_vec_znx(*a_cols, a_size);
782 a_have.raw_mut().copy_from_slice(a_want.raw());
783
784 func_have(&mut a_have);
785
786 // Applies on the minimum matching columns
787 (0..*a_cols).for_each(|i| {
788 // Applies on the minimum matching size
789 (0..a_size).for_each(|j| func_want(a_want.at_poly_mut(i, j)));
790 });
791
792 assert_eq!(a_have.raw(), a_want.raw());
793 });
794 }
795}