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::{apply_binary_op, apply_unary_op, switch_degree, znx_post_process_ternary_op, Backend, Module, VecZnx, ZnxBase, ZnxBasics, ZnxInfos, ZnxLayout};
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 apply_binary_op::<B, VecZnx, 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 apply_binary_op::<B, VecZnx, 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 apply_unary_op::<B, VecZnx>(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 apply_unary_op::<B, VecZnx>(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 apply_unary_op::<B, VecZnx>(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#[cfg(test)]
346mod tests {
347 use crate::{
348 Backend, FFT64, Module, Sampling, VecZnx, VecZnxOps, ZnxBasics, ZnxInfos, ZnxLayout, ffi::vec_znx,
349 znx_post_process_ternary_op,
350 };
351 use itertools::izip;
352 use sampling::source::Source;
353 use std::cmp::min;
354
355 #[test]
356 fn vec_znx_add() {
357 let n: usize = 8;
358 let module: Module<FFT64> = Module::<FFT64>::new(n);
359 let op = |cv: &mut [i64], av: &[i64], bv: &[i64]| {
360 izip!(cv.iter_mut(), bv.iter(), av.iter()).for_each(|(ci, bi, ai)| *ci = *bi + *ai);
361 };
362 test_binary_op::<false, _>(
363 &module,
364 &|c: &mut VecZnx, a: &VecZnx, b: &VecZnx| module.vec_znx_add(c, a, b),
365 op,
366 );
367 }
368
369 #[test]
370 fn vec_znx_add_inplace() {
371 let n: usize = 8;
372 let module: Module<FFT64> = Module::<FFT64>::new(n);
373 let op = |bv: &mut [i64], av: &[i64]| {
374 izip!(bv.iter_mut(), av.iter()).for_each(|(bi, ai)| *bi = *bi + *ai);
375 };
376 test_binary_op_inplace::<false, _>(
377 &module,
378 &|b: &mut VecZnx, a: &VecZnx| module.vec_znx_add_inplace(b, a),
379 op,
380 );
381 }
382
383 #[test]
384 fn vec_znx_sub() {
385 let n: usize = 8;
386 let module: Module<FFT64> = Module::<FFT64>::new(n);
387 let op = |cv: &mut [i64], av: &[i64], bv: &[i64]| {
388 izip!(cv.iter_mut(), bv.iter(), av.iter()).for_each(|(ci, bi, ai)| *ci = *bi - *ai);
389 };
390 test_binary_op::<true, _>(
391 &module,
392 &|c: &mut VecZnx, a: &VecZnx, b: &VecZnx| module.vec_znx_sub(c, a, b),
393 op,
394 );
395 }
396
397 #[test]
398 fn vec_znx_sub_ab_inplace() {
399 let n: usize = 8;
400 let module: Module<FFT64> = Module::<FFT64>::new(n);
401 let op = |bv: &mut [i64], av: &[i64]| {
402 izip!(bv.iter_mut(), av.iter()).for_each(|(bi, ai)| *bi = *ai - *bi);
403 };
404 test_binary_op_inplace::<true, _>(
405 &module,
406 &|b: &mut VecZnx, a: &VecZnx| module.vec_znx_sub_ab_inplace(b, a),
407 op,
408 );
409 }
410
411 #[test]
412 fn vec_znx_sub_ba_inplace() {
413 let n: usize = 8;
414 let module: Module<FFT64> = Module::<FFT64>::new(n);
415 let op = |bv: &mut [i64], av: &[i64]| {
416 izip!(bv.iter_mut(), av.iter()).for_each(|(bi, ai)| *bi = *bi - *ai);
417 };
418 test_binary_op_inplace::<false, _>(
419 &module,
420 &|b: &mut VecZnx, a: &VecZnx| module.vec_znx_sub_ba_inplace(b, a),
421 op,
422 );
423 }
424
425 #[test]
426 fn vec_znx_negate() {
427 let n: usize = 8;
428 let module: Module<FFT64> = Module::<FFT64>::new(n);
429 let op = |b: &mut [i64], a: &[i64]| {
430 izip!(b.iter_mut(), a.iter()).for_each(|(bi, ai)| *bi = -*ai);
431 };
432 test_unary_op(
433 &module,
434 |b: &mut VecZnx, a: &VecZnx| module.vec_znx_negate(b, a),
435 op,
436 )
437 }
438
439 #[test]
440 fn vec_znx_negate_inplace() {
441 let n: usize = 8;
442 let module: Module<FFT64> = Module::<FFT64>::new(n);
443 let op = |a: &mut [i64]| a.iter_mut().for_each(|xi| *xi = -*xi);
444 test_unary_op_inplace(
445 &module,
446 |a: &mut VecZnx| module.vec_znx_negate_inplace(a),
447 op,
448 )
449 }
450
451 #[test]
452 fn vec_znx_rotate() {
453 let n: usize = 8;
454 let module: Module<FFT64> = Module::<FFT64>::new(n);
455 let k: i64 = 53;
456 let op = |b: &mut [i64], a: &[i64]| {
457 assert_eq!(b.len(), a.len());
458 b.copy_from_slice(a);
459
460 let mut k_mod2n: i64 = k % (2 * n as i64);
461 if k_mod2n < 0 {
462 k_mod2n += 2 * n as i64;
463 }
464 let sign: i64 = (k_mod2n.abs() / (n as i64)) & 1;
465 let k_modn: i64 = k_mod2n % (n as i64);
466
467 b.rotate_right(k_modn as usize);
468 b[0..k_modn as usize].iter_mut().for_each(|x| *x = -*x);
469
470 if sign == 1 {
471 b.iter_mut().for_each(|x| *x = -*x);
472 }
473 };
474 test_unary_op(
475 &module,
476 |b: &mut VecZnx, a: &VecZnx| module.vec_znx_rotate(k, b, a),
477 op,
478 )
479 }
480
481 #[test]
482 fn vec_znx_rotate_inplace() {
483 let n: usize = 8;
484 let module: Module<FFT64> = Module::<FFT64>::new(n);
485 let k: i64 = 53;
486 let rot = |a: &mut [i64]| {
487 let mut k_mod2n: i64 = k % (2 * n as i64);
488 if k_mod2n < 0 {
489 k_mod2n += 2 * n as i64;
490 }
491 let sign: i64 = (k_mod2n.abs() / (n as i64)) & 1;
492 let k_modn: i64 = k_mod2n % (n as i64);
493
494 a.rotate_right(k_modn as usize);
495 a[0..k_modn as usize].iter_mut().for_each(|x| *x = -*x);
496
497 if sign == 1 {
498 a.iter_mut().for_each(|x| *x = -*x);
499 }
500 };
501 test_unary_op_inplace(
502 &module,
503 |a: &mut VecZnx| module.vec_znx_rotate_inplace(k, a),
504 rot,
505 )
506 }
507
508 #[test]
509 fn vec_znx_automorphism() {
510 let n: usize = 8;
511 let module: Module<FFT64> = Module::<FFT64>::new(n);
512 let k: i64 = -5;
513 let op = |b: &mut [i64], a: &[i64]| {
514 assert_eq!(b.len(), a.len());
515 unsafe {
516 vec_znx::vec_znx_automorphism(
517 module.ptr,
518 k,
519 b.as_mut_ptr(),
520 1u64,
521 n as u64,
522 a.as_ptr(),
523 1u64,
524 n as u64,
525 );
526 }
527 };
528 test_unary_op(
529 &module,
530 |b: &mut VecZnx, a: &VecZnx| module.vec_znx_automorphism(k, b, a),
531 op,
532 )
533 }
534
535 #[test]
536 fn vec_znx_automorphism_inplace() {
537 let n: usize = 8;
538 let module: Module<FFT64> = Module::<FFT64>::new(n);
539 let k: i64 = -5;
540 let op = |a: &mut [i64]| unsafe {
541 vec_znx::vec_znx_automorphism(
542 module.ptr,
543 k,
544 a.as_mut_ptr(),
545 1u64,
546 n as u64,
547 a.as_ptr(),
548 1u64,
549 n as u64,
550 );
551 };
552 test_unary_op_inplace(
553 &module,
554 |a: &mut VecZnx| module.vec_znx_automorphism_inplace(k, a),
555 op,
556 )
557 }
558
559 fn test_binary_op<const NEGATE: bool, B: Backend>(
560 module: &Module<B>,
561 func_have: impl Fn(&mut VecZnx, &VecZnx, &VecZnx),
562 func_want: impl Fn(&mut [i64], &[i64], &[i64]),
563 ) {
564 let a_size: usize = 3;
565 let b_size: usize = 4;
566 let c_size: usize = 5;
567 let mut source: Source = Source::new([0u8; 32]);
568
569 [1usize, 2, 3].iter().for_each(|a_cols| {
570 [1usize, 2, 3].iter().for_each(|b_cols| {
571 [1usize, 2, 3].iter().for_each(|c_cols| {
572 let min_ab_cols: usize = min(*a_cols, *b_cols);
573 let min_cols: usize = min(*c_cols, min_ab_cols);
574 let min_size: usize = min(c_size, min(a_size, b_size));
575
576 // Allocats a and populates with random values.
577 let mut a: VecZnx = module.new_vec_znx(*a_cols, a_size);
578 (0..*a_cols).for_each(|i| {
579 module.fill_uniform(3, &mut a, i, a_size, &mut source);
580 });
581
582 // Allocats b and populates with random values.
583 let mut b: VecZnx = module.new_vec_znx(*b_cols, b_size);
584 (0..*b_cols).for_each(|i| {
585 module.fill_uniform(3, &mut b, i, b_size, &mut source);
586 });
587
588 // Allocats c and populates with random values.
589 let mut c_have: VecZnx = module.new_vec_znx(*c_cols, c_size);
590 (0..c_have.cols()).for_each(|i| {
591 module.fill_uniform(3, &mut c_have, i, c_size, &mut source);
592 });
593
594 // Applies the function to test
595 func_have(&mut c_have, &a, &b);
596
597 let mut c_want: VecZnx = module.new_vec_znx(*c_cols, c_size);
598
599 // Applies the reference function and expected behavior.
600 // Adds with the minimum matching columns
601 (0..min_cols).for_each(|i| {
602 // Adds with th eminimum matching size
603 (0..min_size).for_each(|j| {
604 func_want(c_want.at_poly_mut(i, j), b.at_poly(i, j), a.at_poly(i, j));
605 });
606
607 if a_size > b_size {
608 // Copies remaining size of lh if lh.size() > rh.size()
609 (min_size..a_size).for_each(|j| {
610 izip!(c_want.at_poly_mut(i, j).iter_mut(), a.at_poly(i, j).iter()).for_each(|(ci, ai)| *ci = *ai);
611 if NEGATE {
612 c_want.at_poly_mut(i, j).iter_mut().for_each(|x| *x = -*x);
613 }
614 });
615 } else {
616 // Copies the remaining size of rh if the are greater
617 (min_size..b_size).for_each(|j| {
618 izip!(c_want.at_poly_mut(i, j).iter_mut(), b.at_poly(i, j).iter()).for_each(|(ci, bi)| *ci = *bi);
619 if NEGATE {
620 c_want.at_poly_mut(i, j).iter_mut().for_each(|x| *x = -*x);
621 }
622 });
623 }
624 });
625
626 znx_post_process_ternary_op::<_, NEGATE>(&mut c_want, &a, &b);
627
628 assert_eq!(c_have.raw(), c_want.raw());
629 });
630 });
631 });
632 }
633
634 fn test_binary_op_inplace<const NEGATE: bool, B: Backend>(
635 module: &Module<B>,
636 func_have: impl Fn(&mut VecZnx, &VecZnx),
637 func_want: impl Fn(&mut [i64], &[i64]),
638 ) {
639 let a_size: usize = 3;
640 let b_size: usize = 5;
641 let mut source = Source::new([0u8; 32]);
642
643 [1usize, 2, 3].iter().for_each(|a_cols| {
644 [1usize, 2, 3].iter().for_each(|b_cols| {
645 let min_cols: usize = min(*b_cols, *a_cols);
646 let min_size: usize = min(b_size, a_size);
647
648 // Allocats a and populates with random values.
649 let mut a: VecZnx = module.new_vec_znx(*a_cols, a_size);
650 (0..*a_cols).for_each(|i| {
651 module.fill_uniform(3, &mut a, i, a_size, &mut source);
652 });
653
654 // Allocats b and populates with random values.
655 let mut b_have: VecZnx = module.new_vec_znx(*b_cols, b_size);
656 (0..*b_cols).for_each(|i| {
657 module.fill_uniform(3, &mut b_have, i, b_size, &mut source);
658 });
659
660 let mut b_want: VecZnx = module.new_vec_znx(*b_cols, b_size);
661 b_want.raw_mut().copy_from_slice(b_have.raw());
662
663 // Applies the function to test.
664 func_have(&mut b_have, &a);
665
666 // Applies the reference function and expected behavior.
667 // Applies with the minimum matching columns
668 (0..min_cols).for_each(|i| {
669 // Adds with th eminimum matching size
670 (0..min_size).for_each(|j| func_want(b_want.at_poly_mut(i, j), a.at_poly(i, j)));
671 if NEGATE {
672 (min_size..b_size).for_each(|j| {
673 b_want.at_poly_mut(i, j).iter_mut().for_each(|x| *x = -*x);
674 });
675 }
676 });
677
678 assert_eq!(b_have.raw(), b_want.raw());
679 });
680 });
681 }
682
683 fn test_unary_op<B: Backend>(
684 module: &Module<B>,
685 func_have: impl Fn(&mut VecZnx, &VecZnx),
686 func_want: impl Fn(&mut [i64], &[i64]),
687 ) {
688 let a_size: usize = 3;
689 let b_size: usize = 5;
690 let mut source = Source::new([0u8; 32]);
691
692 [1usize, 2, 3].iter().for_each(|a_cols| {
693 [1usize, 2, 3].iter().for_each(|b_cols| {
694 let min_cols: usize = min(*b_cols, *a_cols);
695 let min_size: usize = min(b_size, a_size);
696
697 // Allocats a and populates with random values.
698 let mut a: VecZnx = module.new_vec_znx(*a_cols, a_size);
699 (0..a.cols()).for_each(|i| {
700 module.fill_uniform(3, &mut a, i, a_size, &mut source);
701 });
702
703 // Allocats b and populates with random values.
704 let mut b_have: VecZnx = module.new_vec_znx(*b_cols, b_size);
705 (0..b_have.cols()).for_each(|i| {
706 module.fill_uniform(3, &mut b_have, i, b_size, &mut source);
707 });
708
709 let mut b_want: VecZnx = module.new_vec_znx(*b_cols, b_size);
710
711 // Applies the function to test.
712 func_have(&mut b_have, &a);
713
714 // Applies the reference function and expected behavior.
715 // Applies on the minimum matching columns
716 (0..min_cols).for_each(|i| {
717 // Applies on the minimum matching size
718 (0..min_size).for_each(|j| func_want(b_want.at_poly_mut(i, j), a.at_poly(i, j)));
719
720 // Zeroes the unmatching size
721 (min_size..b_size).for_each(|j| {
722 b_want.zero_at(i, j);
723 })
724 });
725
726 // Zeroes the unmatching columns
727 (min_cols..*b_cols).for_each(|i| {
728 (0..b_size).for_each(|j| {
729 b_want.zero_at(i, j);
730 })
731 });
732
733 assert_eq!(b_have.raw(), b_want.raw());
734 });
735 });
736 }
737
738 fn test_unary_op_inplace<B: Backend>(module: &Module<B>, func_have: impl Fn(&mut VecZnx), func_want: impl Fn(&mut [i64])) {
739 let a_size: usize = 3;
740 let mut source = Source::new([0u8; 32]);
741 [1usize, 2, 3].iter().for_each(|a_cols| {
742 let mut a_have: VecZnx = module.new_vec_znx(*a_cols, a_size);
743 (0..*a_cols).for_each(|i| {
744 module.fill_uniform(3, &mut a_have, i, a_size, &mut source);
745 });
746
747 // Allocats a and populates with random values.
748 let mut a_want: VecZnx = module.new_vec_znx(*a_cols, a_size);
749 a_have.raw_mut().copy_from_slice(a_want.raw());
750
751 // Applies the function to test.
752 func_have(&mut a_have);
753
754 // Applies the reference function and expected behavior.
755 // Applies on the minimum matching columns
756 (0..*a_cols).for_each(|i| {
757 // Applies on the minimum matching size
758 (0..a_size).for_each(|j| func_want(a_want.at_poly_mut(i, j)));
759 });
760
761 assert_eq!(a_have.raw(), a_want.raw());
762 });
763 }
764}