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