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