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