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