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 small polynomials per column.
12 fn new_vec_znx(&self, cols: usize, size: usize) -> VecZnx;
13
14 /// Instantiates a new [VecZnx] from a slice of bytes.
15 /// The returned [VecZnx] takes ownership of the slice of bytes.
16 ///
17 /// # Arguments
18 ///
19 /// * `cols`: the number of polynomials.
20 /// * `size`: the number small polynomials per column.
21 ///
22 /// # Panic
23 /// Requires the slice of bytes to be equal to [VecZnxOps::bytes_of_vec_znx].
24 fn new_vec_znx_from_bytes(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnx;
25
26 /// Instantiates a new [VecZnx] from a slice of bytes.
27 /// The returned [VecZnx] does take ownership of the slice of bytes.
28 ///
29 /// # Arguments
30 ///
31 /// * `cols`: the number of polynomials.
32 /// * `size`: the number small polynomials per column.
33 ///
34 /// # Panic
35 /// Requires the slice of bytes to be equal to [VecZnxOps::bytes_of_vec_znx].
36 fn new_vec_znx_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnx;
37
38 /// Returns the number of bytes necessary to allocate
39 /// a new [VecZnx] through [VecZnxOps::new_vec_znx_from_bytes]
40 /// or [VecZnxOps::new_vec_znx_from_bytes_borrow].
41 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize;
42
43 /// Returns the minimum number of bytes necessary for normalization.
44 fn vec_znx_normalize_tmp_bytes(&self, cols: usize) -> usize;
45
46 /// Adds `a` to `b` and write the result on `c`.
47 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
48
49 /// Adds `a` to `b` and write the result on `b`.
50 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx);
51
52 /// Subtracts `b` to `a` and write the result on `c`.
53 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
54
55 /// Subtracts `a` to `b` and write the result on `b`.
56 fn vec_znx_sub_ab_inplace(&self, b: &mut VecZnx, a: &VecZnx);
57
58 /// Subtracts `b` to `a` and write the result on `b`.
59 fn vec_znx_sub_ba_inplace(&self, b: &mut VecZnx, a: &VecZnx);
60
61 // Negates `a` and stores the result on `b`.
62 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx);
63
64 /// Negages `a` and stores the result on `a`.
65 fn vec_znx_negate_inplace(&self, a: &mut VecZnx);
66
67 /// Multiplies `a` by X^k and stores the result on `b`.
68 fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
69
70 /// Multiplies `a` by X^k and stores the result on `a`.
71 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx);
72
73 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `b`.
74 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
75
76 /// Applies the automorphism X^i -> X^ik on `a` and stores the result on `a`.
77 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx);
78
79 /// Splits b into subrings and copies them them into a.
80 ///
81 /// # Panics
82 ///
83 /// This method requires that all [VecZnx] of b have the same ring degree
84 /// and that b.n() * b.len() <= a.n()
85 fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx);
86
87 /// Merges the subrings a into b.
88 ///
89 /// # Panics
90 ///
91 /// This method requires that all [VecZnx] of a have the same ring degree
92 /// and that a.n() * a.len() <= b.n()
93 fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>);
94}
95
96impl<B: Backend> VecZnxOps for Module<B> {
97 fn new_vec_znx(&self, cols: usize, size: usize) -> VecZnx {
98 VecZnx::new(self, cols, size)
99 }
100
101 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize {
102 VecZnx::bytes_of(self, cols, size)
103 }
104
105 fn new_vec_znx_from_bytes(&self, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnx {
106 VecZnx::from_bytes(self, cols, size, bytes)
107 }
108
109 fn new_vec_znx_from_bytes_borrow(&self, cols: usize, size: usize, tmp_bytes: &mut [u8]) -> VecZnx {
110 VecZnx::from_bytes_borrow(self, cols, size, tmp_bytes)
111 }
112
113 fn vec_znx_normalize_tmp_bytes(&self, cols: usize) -> usize {
114 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(self.ptr) as usize * cols }
115 }
116
117 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
118 let op = ffi_ternary_op_factory(
119 self.ptr,
120 c.size(),
121 c.sl(),
122 a.size(),
123 a.sl(),
124 b.size(),
125 b.sl(),
126 vec_znx::vec_znx_add,
127 );
128 vec_znx_apply_binary_op::<B, false>(self, c, a, b, op);
129 }
130
131 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
132 unsafe {
133 let b_ptr: *mut VecZnx = b as *mut VecZnx;
134 Self::vec_znx_add(self, &mut *b_ptr, a, &*b_ptr);
135 }
136 }
137
138 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
139 let op = ffi_ternary_op_factory(
140 self.ptr,
141 c.size(),
142 c.sl(),
143 a.size(),
144 a.sl(),
145 b.size(),
146 b.sl(),
147 vec_znx::vec_znx_sub,
148 );
149 vec_znx_apply_binary_op::<B, true>(self, c, a, b, op);
150 }
151
152 fn vec_znx_sub_ab_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
153 unsafe {
154 let b_ptr: *mut VecZnx = b as *mut VecZnx;
155 Self::vec_znx_sub(self, &mut *b_ptr, a, &*b_ptr);
156 }
157 }
158
159 fn vec_znx_sub_ba_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
160 unsafe {
161 let b_ptr: *mut VecZnx = b as *mut VecZnx;
162 Self::vec_znx_sub(self, &mut *b_ptr, &*b_ptr, a);
163 }
164 }
165
166 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx) {
167 let op = ffi_binary_op_factory_type_0(
168 self.ptr,
169 b.size(),
170 b.sl(),
171 a.size(),
172 a.sl(),
173 vec_znx::vec_znx_negate,
174 );
175 vec_znx_apply_unary_op::<B>(self, b, a, op);
176 }
177
178 fn vec_znx_negate_inplace(&self, a: &mut VecZnx) {
179 unsafe {
180 let a_ptr: *mut VecZnx = a as *mut VecZnx;
181 Self::vec_znx_negate(self, &mut *a_ptr, &*a_ptr);
182 }
183 }
184
185 fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
186 let op = ffi_binary_op_factory_type_1(
187 self.ptr,
188 k,
189 b.size(),
190 b.sl(),
191 a.size(),
192 a.sl(),
193 vec_znx::vec_znx_rotate,
194 );
195 vec_znx_apply_unary_op::<B>(self, b, a, op);
196 }
197
198 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx) {
199 unsafe {
200 let a_ptr: *mut VecZnx = a as *mut VecZnx;
201 Self::vec_znx_rotate(self, k, &mut *a_ptr, &*a_ptr);
202 }
203 }
204
205 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
206 let op = ffi_binary_op_factory_type_1(
207 self.ptr,
208 k,
209 b.size(),
210 b.sl(),
211 a.size(),
212 a.sl(),
213 vec_znx::vec_znx_automorphism,
214 );
215 vec_znx_apply_unary_op::<B>(self, b, a, op);
216 }
217
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 let a_cols: usize = a.cols();
361 let b_cols: usize = b.cols();
362 let c_cols: usize = c.cols();
363 let min_ab_cols: usize = min(a_cols, b_cols);
364 let min_cols: usize = min(c_cols, min_ab_cols);
365 // Applies over shared cols between (a, b, c)
366 (0..min_cols).for_each(|i| op(c.at_poly_mut(i, 0), a.at_poly(i, 0), b.at_poly(i, 0)));
367 // Copies/Negates/Zeroes the remaining cols if op is not inplace.
368 if c.as_ptr() != a.as_ptr() && c.as_ptr() != b.as_ptr() {
369 znx_post_process_ternary_op::<VecZnx, NEGATE>(c, a, b);
370 }
371}
372
373#[inline(always)]
374pub fn vec_znx_apply_unary_op<B: Backend>(module: &Module<B>, b: &mut VecZnx, a: &VecZnx, op: impl Fn(&mut [i64], &[i64])) {
375 #[cfg(debug_assertions)]
376 {
377 assert_eq!(a.n(), module.n());
378 assert_eq!(b.n(), module.n());
379 }
380 let a_cols: usize = a.cols();
381 let b_cols: usize = b.cols();
382 let min_cols: usize = min(a_cols, b_cols);
383 // Applies over the shared cols between (a, b)
384 (0..min_cols).for_each(|i| op(b.at_poly_mut(i, 0), a.at_poly(i, 0)));
385 // Zeroes the remaining cols of b.
386 (min_cols..b_cols).for_each(|i| (0..b.size()).for_each(|j| b.zero_at(i, j)));
387}
388
389#[cfg(test)]
390mod tests {
391 use crate::{
392 Backend, FFT64, Module, Sampling, VecZnx, VecZnxOps, ZnxBasics, ZnxInfos, ZnxLayout, ffi::vec_znx,
393 znx_post_process_ternary_op,
394 };
395 use itertools::izip;
396 use sampling::source::Source;
397 use std::cmp::min;
398
399 #[test]
400 fn vec_znx_add() {
401 let n: usize = 8;
402 let module: Module<FFT64> = Module::<FFT64>::new(n);
403 let op = |cv: &mut [i64], av: &[i64], bv: &[i64]| {
404 izip!(cv.iter_mut(), bv.iter(), av.iter()).for_each(|(ci, bi, ai)| *ci = *bi + *ai);
405 };
406 test_binary_op::<false, _>(
407 &module,
408 &|c: &mut VecZnx, a: &VecZnx, b: &VecZnx| module.vec_znx_add(c, a, b),
409 op,
410 );
411 }
412
413 #[test]
414 fn vec_znx_add_inplace() {
415 let n: usize = 8;
416 let module: Module<FFT64> = Module::<FFT64>::new(n);
417 let op = |bv: &mut [i64], av: &[i64]| {
418 izip!(bv.iter_mut(), av.iter()).for_each(|(bi, ai)| *bi = *bi + *ai);
419 };
420 test_binary_op_inplace::<false, _>(
421 &module,
422 &|b: &mut VecZnx, a: &VecZnx| module.vec_znx_add_inplace(b, a),
423 op,
424 );
425 }
426
427 #[test]
428 fn vec_znx_sub() {
429 let n: usize = 8;
430 let module: Module<FFT64> = Module::<FFT64>::new(n);
431 let op = |cv: &mut [i64], av: &[i64], bv: &[i64]| {
432 izip!(cv.iter_mut(), bv.iter(), av.iter()).for_each(|(ci, bi, ai)| *ci = *bi - *ai);
433 };
434 test_binary_op::<true, _>(
435 &module,
436 &|c: &mut VecZnx, a: &VecZnx, b: &VecZnx| module.vec_znx_sub(c, a, b),
437 op,
438 );
439 }
440
441 #[test]
442 fn vec_znx_sub_ab_inplace() {
443 let n: usize = 8;
444 let module: Module<FFT64> = Module::<FFT64>::new(n);
445 let op = |bv: &mut [i64], av: &[i64]| {
446 izip!(bv.iter_mut(), av.iter()).for_each(|(bi, ai)| *bi = *ai - *bi);
447 };
448 test_binary_op_inplace::<true, _>(
449 &module,
450 &|b: &mut VecZnx, a: &VecZnx| module.vec_znx_sub_ab_inplace(b, a),
451 op,
452 );
453 }
454
455 #[test]
456 fn vec_znx_sub_ba_inplace() {
457 let n: usize = 8;
458 let module: Module<FFT64> = Module::<FFT64>::new(n);
459 let op = |bv: &mut [i64], av: &[i64]| {
460 izip!(bv.iter_mut(), av.iter()).for_each(|(bi, ai)| *bi = *bi - *ai);
461 };
462 test_binary_op_inplace::<false, _>(
463 &module,
464 &|b: &mut VecZnx, a: &VecZnx| module.vec_znx_sub_ba_inplace(b, a),
465 op,
466 );
467 }
468
469 #[test]
470 fn vec_znx_negate() {
471 let n: usize = 8;
472 let module: Module<FFT64> = Module::<FFT64>::new(n);
473 let op = |b: &mut [i64], a: &[i64]| {
474 izip!(b.iter_mut(), a.iter()).for_each(|(bi, ai)| *bi = -*ai);
475 };
476 test_unary_op(
477 &module,
478 |b: &mut VecZnx, a: &VecZnx| module.vec_znx_negate(b, a),
479 op,
480 )
481 }
482
483 #[test]
484 fn vec_znx_negate_inplace() {
485 let n: usize = 8;
486 let module: Module<FFT64> = Module::<FFT64>::new(n);
487 let op = |a: &mut [i64]| a.iter_mut().for_each(|xi| *xi = -*xi);
488 test_unary_op_inplace(
489 &module,
490 |a: &mut VecZnx| module.vec_znx_negate_inplace(a),
491 op,
492 )
493 }
494
495 #[test]
496 fn vec_znx_rotate() {
497 let n: usize = 8;
498 let module: Module<FFT64> = Module::<FFT64>::new(n);
499 let k: i64 = 53;
500 let op = |b: &mut [i64], a: &[i64]| {
501 assert_eq!(b.len(), a.len());
502 b.copy_from_slice(a);
503
504 let mut k_mod2n: i64 = k % (2 * n as i64);
505 if k_mod2n < 0 {
506 k_mod2n += 2 * n as i64;
507 }
508 let sign: i64 = (k_mod2n.abs() / (n as i64)) & 1;
509 let k_modn: i64 = k_mod2n % (n as i64);
510
511 b.rotate_right(k_modn as usize);
512 b[0..k_modn as usize].iter_mut().for_each(|x| *x = -*x);
513
514 if sign == 1 {
515 b.iter_mut().for_each(|x| *x = -*x);
516 }
517 };
518 test_unary_op(
519 &module,
520 |b: &mut VecZnx, a: &VecZnx| module.vec_znx_rotate(k, b, a),
521 op,
522 )
523 }
524
525 #[test]
526 fn vec_znx_rotate_inplace() {
527 let n: usize = 8;
528 let module: Module<FFT64> = Module::<FFT64>::new(n);
529 let k: i64 = 53;
530 let rot = |a: &mut [i64]| {
531 let mut k_mod2n: i64 = k % (2 * n as i64);
532 if k_mod2n < 0 {
533 k_mod2n += 2 * n as i64;
534 }
535 let sign: i64 = (k_mod2n.abs() / (n as i64)) & 1;
536 let k_modn: i64 = k_mod2n % (n as i64);
537
538 a.rotate_right(k_modn as usize);
539 a[0..k_modn as usize].iter_mut().for_each(|x| *x = -*x);
540
541 if sign == 1 {
542 a.iter_mut().for_each(|x| *x = -*x);
543 }
544 };
545 test_unary_op_inplace(
546 &module,
547 |a: &mut VecZnx| module.vec_znx_rotate_inplace(k, a),
548 rot,
549 )
550 }
551
552 #[test]
553 fn vec_znx_automorphism() {
554 let n: usize = 8;
555 let module: Module<FFT64> = Module::<FFT64>::new(n);
556 let k: i64 = -5;
557 let op = |b: &mut [i64], a: &[i64]| {
558 assert_eq!(b.len(), a.len());
559 unsafe {
560 vec_znx::vec_znx_automorphism(
561 module.ptr,
562 k,
563 b.as_mut_ptr(),
564 1u64,
565 n as u64,
566 a.as_ptr(),
567 1u64,
568 n as u64,
569 );
570 }
571 };
572 test_unary_op(
573 &module,
574 |b: &mut VecZnx, a: &VecZnx| module.vec_znx_automorphism(k, b, a),
575 op,
576 )
577 }
578
579 #[test]
580 fn vec_znx_automorphism_inplace() {
581 let n: usize = 8;
582 let module: Module<FFT64> = Module::<FFT64>::new(n);
583 let k: i64 = -5;
584 let op = |a: &mut [i64]| unsafe {
585 vec_znx::vec_znx_automorphism(
586 module.ptr,
587 k,
588 a.as_mut_ptr(),
589 1u64,
590 n as u64,
591 a.as_ptr(),
592 1u64,
593 n as u64,
594 );
595 };
596 test_unary_op_inplace(
597 &module,
598 |a: &mut VecZnx| module.vec_znx_automorphism_inplace(k, a),
599 op,
600 )
601 }
602
603 fn test_binary_op<const NEGATE: bool, B: Backend>(
604 module: &Module<B>,
605 func_have: impl Fn(&mut VecZnx, &VecZnx, &VecZnx),
606 func_want: impl Fn(&mut [i64], &[i64], &[i64]),
607 ) {
608 let a_size: usize = 3;
609 let b_size: usize = 4;
610 let c_size: usize = 5;
611 let mut source: Source = Source::new([0u8; 32]);
612
613 [1usize, 2, 3].iter().for_each(|a_cols| {
614 [1usize, 2, 3].iter().for_each(|b_cols| {
615 [1usize, 2, 3].iter().for_each(|c_cols| {
616 let min_ab_cols: usize = min(*a_cols, *b_cols);
617 let min_cols: usize = min(*c_cols, min_ab_cols);
618 let min_size: usize = min(c_size, min(a_size, b_size));
619
620 // Allocats a and populates with random values.
621 let mut a: VecZnx = module.new_vec_znx(*a_cols, a_size);
622 (0..*a_cols).for_each(|i| {
623 module.fill_uniform(3, &mut a, i, a_size, &mut source);
624 });
625
626 // Allocats b and populates with random values.
627 let mut b: VecZnx = module.new_vec_znx(*b_cols, b_size);
628 (0..*b_cols).for_each(|i| {
629 module.fill_uniform(3, &mut b, i, b_size, &mut source);
630 });
631
632 // Allocats c and populates with random values.
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 // Applies the function to test
639 func_have(&mut c_have, &a, &b);
640
641 let mut c_want: VecZnx = module.new_vec_znx(*c_cols, c_size);
642
643 // Applies the reference function and expected behavior.
644 // Adds with the minimum matching columns
645 (0..min_cols).for_each(|i| {
646 // Adds with th eminimum matching size
647 (0..min_size).for_each(|j| {
648 func_want(c_want.at_poly_mut(i, j), b.at_poly(i, j), a.at_poly(i, j));
649 });
650
651 if a_size > b_size {
652 // Copies remaining size of lh if lh.size() > rh.size()
653 (min_size..a_size).for_each(|j| {
654 izip!(c_want.at_poly_mut(i, j).iter_mut(), a.at_poly(i, j).iter()).for_each(|(ci, ai)| *ci = *ai);
655 if NEGATE {
656 c_want.at_poly_mut(i, j).iter_mut().for_each(|x| *x = -*x);
657 }
658 });
659 } else {
660 // Copies the remaining size of rh if the are greater
661 (min_size..b_size).for_each(|j| {
662 izip!(c_want.at_poly_mut(i, j).iter_mut(), b.at_poly(i, j).iter()).for_each(|(ci, bi)| *ci = *bi);
663 if NEGATE {
664 c_want.at_poly_mut(i, j).iter_mut().for_each(|x| *x = -*x);
665 }
666 });
667 }
668 });
669
670 znx_post_process_ternary_op::<_, NEGATE>(&mut c_want, &a, &b);
671
672 assert_eq!(c_have.raw(), c_want.raw());
673 });
674 });
675 });
676 }
677
678 fn test_binary_op_inplace<const NEGATE: bool, B: Backend>(
679 module: &Module<B>,
680 func_have: impl Fn(&mut VecZnx, &VecZnx),
681 func_want: impl Fn(&mut [i64], &[i64]),
682 ) {
683 let a_size: usize = 3;
684 let b_size: usize = 5;
685 let mut source = Source::new([0u8; 32]);
686
687 [1usize, 2, 3].iter().for_each(|a_cols| {
688 [1usize, 2, 3].iter().for_each(|b_cols| {
689 let min_cols: usize = min(*b_cols, *a_cols);
690 let min_size: usize = min(b_size, a_size);
691
692 // Allocats a and populates with random values.
693 let mut a: VecZnx = module.new_vec_znx(*a_cols, a_size);
694 (0..*a_cols).for_each(|i| {
695 module.fill_uniform(3, &mut a, i, a_size, &mut source);
696 });
697
698 // Allocats b and populates with random values.
699 let mut b_have: VecZnx = module.new_vec_znx(*b_cols, b_size);
700 (0..*b_cols).for_each(|i| {
701 module.fill_uniform(3, &mut b_have, i, b_size, &mut source);
702 });
703
704 let mut b_want: VecZnx = module.new_vec_znx(*b_cols, b_size);
705 b_want.raw_mut().copy_from_slice(b_have.raw());
706
707 // Applies the function to test.
708 func_have(&mut b_have, &a);
709
710 // Applies the reference function and expected behavior.
711 // Applies with the minimum matching columns
712 (0..min_cols).for_each(|i| {
713 // Adds with th eminimum matching size
714 (0..min_size).for_each(|j| func_want(b_want.at_poly_mut(i, j), a.at_poly(i, j)));
715 if NEGATE {
716 (min_size..b_size).for_each(|j| {
717 b_want.at_poly_mut(i, j).iter_mut().for_each(|x| *x = -*x);
718 });
719 }
720 });
721
722 assert_eq!(b_have.raw(), b_want.raw());
723 });
724 });
725 }
726
727 fn test_unary_op<B: Backend>(
728 module: &Module<B>,
729 func_have: impl Fn(&mut VecZnx, &VecZnx),
730 func_want: impl Fn(&mut [i64], &[i64]),
731 ) {
732 let a_size: usize = 3;
733 let b_size: usize = 5;
734 let mut source = Source::new([0u8; 32]);
735
736 [1usize, 2, 3].iter().for_each(|a_cols| {
737 [1usize, 2, 3].iter().for_each(|b_cols| {
738 let min_cols: usize = min(*b_cols, *a_cols);
739 let min_size: usize = min(b_size, a_size);
740
741 // Allocats a and populates with random values.
742 let mut a: VecZnx = module.new_vec_znx(*a_cols, a_size);
743 (0..a.cols()).for_each(|i| {
744 module.fill_uniform(3, &mut a, i, a_size, &mut source);
745 });
746
747 // Allocats b and populates with random values.
748 let mut b_have: VecZnx = module.new_vec_znx(*b_cols, b_size);
749 (0..b_have.cols()).for_each(|i| {
750 module.fill_uniform(3, &mut b_have, i, b_size, &mut source);
751 });
752
753 let mut b_want: VecZnx = module.new_vec_znx(*b_cols, b_size);
754
755 // Applies the function to test.
756 func_have(&mut b_have, &a);
757
758 // Applies the reference function and expected behavior.
759 // Applies on the minimum matching columns
760 (0..min_cols).for_each(|i| {
761 // Applies on the minimum matching size
762 (0..min_size).for_each(|j| func_want(b_want.at_poly_mut(i, j), a.at_poly(i, j)));
763
764 // Zeroes the unmatching size
765 (min_size..b_size).for_each(|j| {
766 b_want.zero_at(i, j);
767 })
768 });
769
770 // Zeroes the unmatching columns
771 (min_cols..*b_cols).for_each(|i| {
772 (0..b_size).for_each(|j| {
773 b_want.zero_at(i, j);
774 })
775 });
776
777 assert_eq!(b_have.raw(), b_want.raw());
778 });
779 });
780 }
781
782 fn test_unary_op_inplace<B: Backend>(module: &Module<B>, func_have: impl Fn(&mut VecZnx), func_want: impl Fn(&mut [i64])) {
783 let a_size: usize = 3;
784 let mut source = Source::new([0u8; 32]);
785 [1usize, 2, 3].iter().for_each(|a_cols| {
786 let mut a_have: VecZnx = module.new_vec_znx(*a_cols, a_size);
787 (0..*a_cols).for_each(|i| {
788 module.fill_uniform(3, &mut a_have, i, a_size, &mut source);
789 });
790
791 // Allocats a and populates with random values.
792 let mut a_want: VecZnx = module.new_vec_znx(*a_cols, a_size);
793 a_have.raw_mut().copy_from_slice(a_want.raw());
794
795 // Applies the function to test.
796 func_have(&mut a_have);
797
798 // Applies the reference function and expected behavior.
799 // Applies on the minimum matching columns
800 (0..*a_cols).for_each(|i| {
801 // Applies on the minimum matching size
802 (0..a_size).for_each(|j| func_want(a_want.at_poly_mut(i, j)));
803 });
804
805 assert_eq!(a_have.raw(), a_want.raw());
806 });
807 }
808}