poulpy/base2k/src/vec_znx.rs
base2k/src/vec_znx.rsRaw
1use crate::Backend;
2use crate::cast_mut;
3use crate::ffi::vec_znx;
4use crate::ffi::znx;
5use crate::{Infos, Module};
6use crate::{alloc_aligned, assert_alignement};
7use itertools::izip;
8use std::cmp::min;
9
10/// [VecZnx] represents collection of contiguously stacked vector of small norm polynomials of
11/// Zn\[X\] with [i64] coefficients.
12/// A [VecZnx] is composed of multiple Zn\[X\] polynomials stored in a single contiguous array
13/// in the memory.
14///
15/// # Example
16///
17/// Given 3 polynomials (a, b, c) of Zn\[X\], each with 4 columns, then the memory
18/// layout is: `[a0, b0, c0, a1, b1, c1, a2, b2, c2, a3, b3, c3]`, where ai, bi, ci
19/// are small polynomials of Zn\[X\].
20#[derive(Clone)]
21pub struct VecZnx {
22 /// Polynomial degree.
23 pub n: usize,
24
25 /// The number of polynomials
26 pub cols: usize,
27
28 /// The number of limbs per polynomial (a.k.a small polynomials).
29 pub limbs: usize,
30
31 /// Polynomial coefficients, as a contiguous array. Each col is equally spaced by n.
32 pub data: Vec<i64>,
33
34 /// Pointer to data (data can be enpty if [VecZnx] borrows space instead of owning it).
35 pub ptr: *mut i64,
36}
37
38pub fn bytes_of_vec_znx(n: usize, cols: usize, limbs: usize) -> usize {
39 n * cols * limbs * size_of::<i64>()
40}
41
42impl VecZnx {
43 /// Returns a new struct implementing [VecZnx] with the provided data as backing array.
44 ///
45 /// The struct will take ownership of buf[..[Self::bytes_of]]
46 ///
47 /// User must ensure that data is properly alligned and that
48 /// the limbs of data is equal to [Self::bytes_of].
49 pub fn from_bytes(n: usize, cols: usize, limbs: usize, bytes: &mut [u8]) -> Self {
50 #[cfg(debug_assertions)]
51 {
52 assert!(cols > 0);
53 assert!(limbs > 0);
54 assert_eq!(bytes.len(), Self::bytes_of(n, cols, limbs));
55 assert_alignement(bytes.as_ptr());
56 }
57 unsafe {
58 let bytes_i64: &mut [i64] = cast_mut::<u8, i64>(bytes);
59 let ptr: *mut i64 = bytes_i64.as_mut_ptr();
60 Self {
61 n: n,
62 cols: cols,
63 limbs: limbs,
64 data: Vec::from_raw_parts(ptr, bytes.len(), bytes.len()),
65 ptr: ptr,
66 }
67 }
68 }
69
70 pub fn from_bytes_borrow(n: usize, cols: usize, limbs: usize, bytes: &mut [u8]) -> Self {
71 #[cfg(debug_assertions)]
72 {
73 assert!(cols > 0);
74 assert!(limbs > 0);
75 assert!(bytes.len() >= Self::bytes_of(n, cols, limbs));
76 assert_alignement(bytes.as_ptr());
77 }
78 Self {
79 n: n,
80 cols: cols,
81 limbs: limbs,
82 data: Vec::new(),
83 ptr: bytes.as_mut_ptr() as *mut i64,
84 }
85 }
86
87 pub fn bytes_of(n: usize, cols: usize, limbs: usize) -> usize {
88 bytes_of_vec_znx(n, cols, limbs)
89 }
90
91 pub fn copy_from(&mut self, a: &Self) {
92 copy_vec_znx_from(self, a);
93 }
94
95 pub fn borrowing(&self) -> bool {
96 self.data.len() == 0
97 }
98
99 /// Total limbs is [Self::n()] * [Self::poly_count()].
100 pub fn raw(&self) -> &[i64] {
101 unsafe { std::slice::from_raw_parts(self.ptr, self.n * self.poly_count()) }
102 }
103
104 /// Returns a reference to backend slice of the receiver.
105 /// Total size is [Self::n()] * [Self::poly_count()].
106 pub fn raw_mut(&mut self) -> &mut [i64] {
107 unsafe { std::slice::from_raw_parts_mut(self.ptr, self.n * self.poly_count()) }
108 }
109
110 /// Returns a non-mutable pointer to the backedn slice of the receiver.
111 pub fn as_ptr(&self) -> *const i64 {
112 self.ptr
113 }
114
115 /// Returns a mutable pointer to the backedn slice of the receiver.
116 pub fn as_mut_ptr(&mut self) -> *mut i64 {
117 self.ptr
118 }
119
120 /// Returns a non-mutable pointer starting a the (i, j)-th small poly.
121 pub fn at_ptr(&self, i: usize, j: usize) -> *const i64 {
122 #[cfg(debug_assertions)]
123 {
124 assert!(i < self.cols());
125 assert!(j < self.limbs());
126 }
127 let offset: usize = self.n * (j * self.cols() + i);
128 self.ptr.wrapping_add(offset)
129 }
130
131 /// Returns a non-mutable reference to the i-th limb.
132 /// The returned array is of size [Self::n()] * [Self::cols()].
133 pub fn at_limb(&self, i: usize) -> &[i64] {
134 unsafe { std::slice::from_raw_parts(self.at_ptr(0, i), self.n * self.cols()) }
135 }
136
137 /// Returns a non-mutable reference to the (i, j)-th poly.
138 /// The returned array is of size [Self::n()].
139 pub fn at_poly(&self, i: usize, j: usize) -> &[i64] {
140 unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n) }
141 }
142
143 /// Returns a mutable pointer starting a the (i, j)-th small poly.
144 pub fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut i64 {
145 #[cfg(debug_assertions)]
146 {
147 assert!(i < self.cols());
148 assert!(j < self.limbs());
149 }
150
151 let offset: usize = self.n * (j * self.cols() + i);
152 self.ptr.wrapping_add(offset)
153 }
154
155 /// Returns a mutable reference to the i-th limb.
156 /// The returned array is of size [Self::n()] * [Self::cols()].
157 pub fn at_limb_mut(&mut self, i: usize) -> &mut [i64] {
158 unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(0, i), self.n * self.cols()) }
159 }
160
161 /// Returns a mutable reference to the (i, j)-th poly.
162 /// The returned array is of size [Self::n()].
163 pub fn at_poly_mut(&mut self, i: usize, j: usize) -> &mut [i64] {
164 unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n) }
165 }
166
167 pub fn zero(&mut self) {
168 unsafe { znx::znx_zero_i64_ref((self.n * self.poly_count()) as u64, self.ptr) }
169 }
170
171 pub fn normalize(&mut self, log_base2k: usize, carry: &mut [u8]) {
172 normalize(log_base2k, self, carry)
173 }
174
175 pub fn rsh(&mut self, log_base2k: usize, k: usize, carry: &mut [u8]) {
176 rsh(log_base2k, self, k, carry)
177 }
178
179 pub fn switch_degree(&self, a: &mut Self) {
180 switch_degree(a, self)
181 }
182
183 // Prints the first `n` coefficients of each limb
184 pub fn print(&self, n: usize) {
185 (0..self.limbs()).for_each(|i| println!("{}: {:?}", i, &self.at_limb(i)[..n]))
186 }
187}
188
189impl Infos for VecZnx {
190 fn n(&self) -> usize {
191 self.n
192 }
193
194 fn log_n(&self) -> usize {
195 (usize::BITS - (self.n() - 1).leading_zeros()) as _
196 }
197
198 fn rows(&self) -> usize {
199 1
200 }
201
202 fn cols(&self) -> usize {
203 self.cols
204 }
205
206 fn limbs(&self) -> usize {
207 self.limbs
208 }
209
210 fn poly_count(&self) -> usize {
211 self.cols * self.limbs
212 }
213}
214
215/// Copies the coefficients of `a` on the receiver.
216/// Copy is done with the minimum size matching both backing arrays.
217/// Panics if the cols do not match.
218pub fn copy_vec_znx_from(b: &mut VecZnx, a: &VecZnx) {
219 assert_eq!(b.cols(), a.cols());
220 let data_a: &[i64] = a.raw();
221 let data_b: &mut [i64] = b.raw_mut();
222 let size = min(data_b.len(), data_a.len());
223 data_b[..size].copy_from_slice(&data_a[..size])
224}
225
226impl VecZnx {
227 /// Allocates a new [VecZnx] composed of #size polynomials of Z\[X\].
228 pub fn new(n: usize, cols: usize, limbs: usize) -> Self {
229 #[cfg(debug_assertions)]
230 {
231 assert!(n > 0);
232 assert!(n & (n - 1) == 0);
233 assert!(cols > 0);
234 assert!(limbs > 0);
235 }
236 let mut data: Vec<i64> = alloc_aligned::<i64>(n * cols * limbs);
237 let ptr: *mut i64 = data.as_mut_ptr();
238 Self {
239 n: n,
240 cols: cols,
241 limbs: limbs,
242 data: data,
243 ptr: ptr,
244 }
245 }
246
247 /// Truncates the precision of the [VecZnx] by k bits.
248 ///
249 /// # Arguments
250 ///
251 /// * `log_base2k`: the base two logarithm of the coefficients decomposition.
252 /// * `k`: the number of bits of precision to drop.
253 pub fn trunc_pow2(&mut self, log_base2k: usize, k: usize) {
254 if k == 0 {
255 return;
256 }
257
258 if !self.borrowing() {
259 self.data
260 .truncate(self.n() * self.cols() * (self.limbs() - k / log_base2k));
261 }
262
263 self.limbs -= k / log_base2k;
264
265 let k_rem: usize = k % log_base2k;
266
267 if k_rem != 0 {
268 let mask: i64 = ((1 << (log_base2k - k_rem - 1)) - 1) << k_rem;
269 self.at_limb_mut(self.limbs() - 1)
270 .iter_mut()
271 .for_each(|x: &mut i64| *x &= mask)
272 }
273 }
274}
275
276pub fn switch_degree(b: &mut VecZnx, a: &VecZnx) {
277 let (n_in, n_out) = (a.n(), b.n());
278 let (gap_in, gap_out): (usize, usize);
279
280 if n_in > n_out {
281 (gap_in, gap_out) = (n_in / n_out, 1)
282 } else {
283 (gap_in, gap_out) = (1, n_out / n_in);
284 b.zero();
285 }
286
287 let limbs: usize = min(a.limbs(), b.limbs());
288
289 (0..limbs).for_each(|i| {
290 izip!(
291 a.at_limb(i).iter().step_by(gap_in),
292 b.at_limb_mut(i).iter_mut().step_by(gap_out)
293 )
294 .for_each(|(x_in, x_out)| *x_out = *x_in);
295 });
296}
297
298fn normalize_tmp_bytes(n: usize, limbs: usize) -> usize {
299 n * limbs * std::mem::size_of::<i64>()
300}
301
302fn normalize(log_base2k: usize, a: &mut VecZnx, tmp_bytes: &mut [u8]) {
303 let n: usize = a.n();
304 let cols: usize = a.cols();
305
306 debug_assert!(
307 tmp_bytes.len() >= normalize_tmp_bytes(n, cols),
308 "invalid tmp_bytes: tmp_bytes.len()={} < normalize_tmp_bytes({}, {})",
309 tmp_bytes.len(),
310 n,
311 cols,
312 );
313 #[cfg(debug_assertions)]
314 {
315 assert_alignement(tmp_bytes.as_ptr())
316 }
317
318 let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
319
320 unsafe {
321 znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
322 (0..a.limbs()).rev().for_each(|i| {
323 znx::znx_normalize(
324 (n * cols) as u64,
325 log_base2k as u64,
326 a.at_mut_ptr(0, i),
327 carry_i64.as_mut_ptr(),
328 a.at_mut_ptr(0, i),
329 carry_i64.as_mut_ptr(),
330 )
331 });
332 }
333}
334
335pub fn rsh_tmp_bytes(n: usize, limbs: usize) -> usize {
336 n * limbs * std::mem::size_of::<i64>()
337}
338
339pub fn rsh(log_base2k: usize, a: &mut VecZnx, k: usize, tmp_bytes: &mut [u8]) {
340 let n: usize = a.n();
341 let limbs: usize = a.limbs();
342
343 #[cfg(debug_assertions)]
344 {
345 assert!(
346 tmp_bytes.len() >= rsh_tmp_bytes(n, limbs),
347 "invalid carry: carry.len()/8={} < rsh_tmp_bytes({}, {})",
348 tmp_bytes.len() >> 3,
349 n,
350 limbs,
351 );
352 assert_alignement(tmp_bytes.as_ptr());
353 }
354
355 let limbs: usize = a.limbs();
356 let size_steps: usize = k / log_base2k;
357
358 a.raw_mut().rotate_right(n * limbs * size_steps);
359 unsafe {
360 znx::znx_zero_i64_ref((n * limbs * size_steps) as u64, a.as_mut_ptr());
361 }
362
363 let k_rem = k % log_base2k;
364
365 if k_rem != 0 {
366 let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
367
368 unsafe {
369 znx::znx_zero_i64_ref((n * limbs) as u64, carry_i64.as_mut_ptr());
370 }
371
372 let log_base2k: usize = log_base2k;
373
374 (size_steps..limbs).for_each(|i| {
375 izip!(carry_i64.iter_mut(), a.at_limb_mut(i).iter_mut()).for_each(|(ci, xi)| {
376 *xi += *ci << log_base2k;
377 *ci = get_base_k_carry(*xi, k_rem);
378 *xi = (*xi - *ci) >> k_rem;
379 });
380 })
381 }
382}
383
384#[inline(always)]
385fn get_base_k_carry(x: i64, k: usize) -> i64 {
386 (x << 64 - k) >> (64 - k)
387}
388
389pub trait VecZnxOps {
390 /// Allocates a new [VecZnx].
391 ///
392 /// # Arguments
393 ///
394 /// * `cols`: the number of polynomials.
395 /// * `limbs`: the number of limbs per polynomial (a.k.a small polynomials).
396 fn new_vec_znx(&self, cols: usize, limbs: usize) -> VecZnx;
397
398 /// Returns the minimum number of bytes necessary to allocate
399 /// a new [VecZnx] through [VecZnx::from_bytes].
400 fn bytes_of_vec_znx(&self, cols: usize, size: usize) -> usize;
401
402 fn vec_znx_normalize_tmp_bytes(&self, cols: usize) -> usize;
403
404 /// c <- a + b.
405 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
406
407 /// b <- b + a.
408 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx);
409
410 /// c <- a - b.
411 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
412
413 /// b <- a - b.
414 fn vec_znx_sub_ab_inplace(&self, b: &mut VecZnx, a: &VecZnx);
415
416 /// b <- b - a.
417 fn vec_znx_sub_ba_inplace(&self, b: &mut VecZnx, a: &VecZnx);
418
419 /// b <- -a.
420 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx);
421
422 /// b <- -b.
423 fn vec_znx_negate_inplace(&self, a: &mut VecZnx);
424
425 /// b <- a * X^k (mod X^{n} + 1)
426 fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
427
428 /// a <- a * X^k (mod X^{n} + 1)
429 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx);
430
431 /// b <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
432 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
433
434 /// a <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
435 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx);
436
437 /// Splits b into subrings and copies them them into a.
438 ///
439 /// # Panics
440 ///
441 /// This method requires that all [VecZnx] of b have the same ring degree
442 /// and that b.n() * b.len() <= a.n()
443 fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx);
444
445 /// Merges the subrings a into b.
446 ///
447 /// # Panics
448 ///
449 /// This method requires that all [VecZnx] of a have the same ring degree
450 /// and that a.n() * a.len() <= b.n()
451 fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>);
452}
453
454impl<B: Backend> VecZnxOps for Module<B> {
455 fn new_vec_znx(&self, cols: usize, limbs: usize) -> VecZnx {
456 VecZnx::new(self.n(), cols, limbs)
457 }
458
459 fn bytes_of_vec_znx(&self, cols: usize, limbs: usize) -> usize {
460 bytes_of_vec_znx(self.n(), cols, limbs)
461 }
462
463 fn vec_znx_normalize_tmp_bytes(&self, cols: usize) -> usize {
464 unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(self.ptr) as usize * cols }
465 }
466
467 // c <- a + b
468 fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
469 let n: usize = self.n();
470 #[cfg(debug_assertions)]
471 {
472 assert_eq!(c.n(), n);
473 assert_eq!(a.n(), n);
474 assert_eq!(b.n(), n);
475 }
476 unsafe {
477 vec_znx::vec_znx_add(
478 self.ptr,
479 c.as_mut_ptr(),
480 c.limbs() as u64,
481 (n * c.cols()) as u64,
482 a.as_ptr(),
483 a.limbs() as u64,
484 (n * a.cols()) as u64,
485 b.as_ptr(),
486 b.limbs() as u64,
487 (n * b.cols()) as u64,
488 )
489 }
490 }
491
492 // b <- a + b
493 fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
494 let n: usize = self.n();
495 #[cfg(debug_assertions)]
496 {
497 assert_eq!(a.n(), n);
498 assert_eq!(b.n(), n);
499 }
500 unsafe {
501 vec_znx::vec_znx_add(
502 self.ptr,
503 b.as_mut_ptr(),
504 b.limbs() as u64,
505 (n * b.cols()) as u64,
506 a.as_ptr(),
507 a.limbs() as u64,
508 (n * a.cols()) as u64,
509 b.as_ptr(),
510 b.limbs() as u64,
511 (n * b.cols()) as u64,
512 )
513 }
514 }
515
516 // c <- a + b
517 fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
518 let n: usize = self.n();
519 #[cfg(debug_assertions)]
520 {
521 assert_eq!(c.n(), n);
522 assert_eq!(a.n(), n);
523 assert_eq!(b.n(), n);
524 }
525 unsafe {
526 vec_znx::vec_znx_sub(
527 self.ptr,
528 c.as_mut_ptr(),
529 c.limbs() as u64,
530 (n * c.cols()) as u64,
531 a.as_ptr(),
532 a.limbs() as u64,
533 (n * a.cols()) as u64,
534 b.as_ptr(),
535 b.limbs() as u64,
536 (n * b.cols()) as u64,
537 )
538 }
539 }
540
541 // b <- a - b
542 fn vec_znx_sub_ab_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
543 let n: usize = self.n();
544 #[cfg(debug_assertions)]
545 {
546 assert_eq!(a.n(), n);
547 assert_eq!(b.n(), n);
548 }
549 unsafe {
550 vec_znx::vec_znx_sub(
551 self.ptr,
552 b.as_mut_ptr(),
553 b.limbs() as u64,
554 (n * b.cols()) as u64,
555 a.as_ptr(),
556 a.limbs() as u64,
557 (n * a.cols()) as u64,
558 b.as_ptr(),
559 b.limbs() as u64,
560 (n * b.cols()) as u64,
561 )
562 }
563 }
564
565 // b <- b - a
566 fn vec_znx_sub_ba_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
567 let n: usize = self.n();
568 #[cfg(debug_assertions)]
569 {
570 assert_eq!(a.n(), n);
571 assert_eq!(b.n(), n);
572 }
573 unsafe {
574 vec_znx::vec_znx_sub(
575 self.ptr,
576 b.as_mut_ptr(),
577 b.limbs() as u64,
578 (n * b.cols()) as u64,
579 b.as_ptr(),
580 b.limbs() as u64,
581 (n * b.cols()) as u64,
582 a.as_ptr(),
583 a.limbs() as u64,
584 (n * a.cols()) as u64,
585 )
586 }
587 }
588
589 fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx) {
590 let n: usize = self.n();
591 #[cfg(debug_assertions)]
592 {
593 assert_eq!(a.n(), n);
594 assert_eq!(b.n(), n);
595 }
596 unsafe {
597 vec_znx::vec_znx_negate(
598 self.ptr,
599 b.as_mut_ptr(),
600 b.limbs() as u64,
601 (n * b.cols()) as u64,
602 a.as_ptr(),
603 a.limbs() as u64,
604 (n * a.cols()) as u64,
605 )
606 }
607 }
608
609 fn vec_znx_negate_inplace(&self, a: &mut VecZnx) {
610 let n: usize = self.n();
611 #[cfg(debug_assertions)]
612 {
613 assert_eq!(a.n(), n);
614 }
615 unsafe {
616 vec_znx::vec_znx_negate(
617 self.ptr,
618 a.as_mut_ptr(),
619 a.limbs() as u64,
620 (n * a.cols()) as u64,
621 a.as_ptr(),
622 a.limbs() as u64,
623 (n * a.cols()) as u64,
624 )
625 }
626 }
627
628 fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
629 let n: usize = self.n();
630 #[cfg(debug_assertions)]
631 {
632 assert_eq!(a.n(), n);
633 assert_eq!(b.n(), n);
634 }
635 unsafe {
636 vec_znx::vec_znx_rotate(
637 self.ptr,
638 k,
639 b.as_mut_ptr(),
640 b.limbs() as u64,
641 (n * b.cols()) as u64,
642 a.as_ptr(),
643 a.limbs() as u64,
644 (n * a.cols()) as u64,
645 )
646 }
647 }
648
649 fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx) {
650 let n: usize = self.n();
651 #[cfg(debug_assertions)]
652 {
653 assert_eq!(a.n(), n);
654 }
655 unsafe {
656 vec_znx::vec_znx_rotate(
657 self.ptr,
658 k,
659 a.as_mut_ptr(),
660 a.limbs() as u64,
661 (n * a.cols()) as u64,
662 a.as_ptr(),
663 a.limbs() as u64,
664 (n * a.cols()) as u64,
665 )
666 }
667 }
668
669 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each size.
670 ///
671 /// # Arguments
672 ///
673 /// * `a`: input.
674 /// * `b`: output.
675 /// * `k`: the power to which to map each coefficients.
676 /// * `a_size`: the number of a_size on which to apply the mapping.
677 ///
678 /// # Panics
679 ///
680 /// The method will panic if the argument `a` is greater than `a.limbs()`.
681 fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
682 let n: usize = self.n();
683 #[cfg(debug_assertions)]
684 {
685 assert_eq!(a.n(), n);
686 assert_eq!(b.n(), n);
687 }
688 unsafe {
689 vec_znx::vec_znx_automorphism(
690 self.ptr,
691 k,
692 b.as_mut_ptr(),
693 b.limbs() as u64,
694 (n * b.cols()) as u64,
695 a.as_ptr(),
696 a.limbs() as u64,
697 (n * a.cols()) as u64,
698 );
699 }
700 }
701
702 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each size.
703 ///
704 /// # Arguments
705 ///
706 /// * `a`: input and output.
707 /// * `k`: the power to which to map each coefficients.
708 /// * `a_size`: the number of size on which to apply the mapping.
709 ///
710 /// # Panics
711 ///
712 /// The method will panic if the argument `size` is greater than `self.limbs()`.
713 fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx) {
714 let n: usize = self.n();
715 #[cfg(debug_assertions)]
716 {
717 assert_eq!(a.n(), n);
718 }
719 unsafe {
720 vec_znx::vec_znx_automorphism(
721 self.ptr,
722 k,
723 a.as_mut_ptr(),
724 a.limbs() as u64,
725 (n * a.cols()) as u64,
726 a.as_ptr(),
727 a.limbs() as u64,
728 (n * a.cols()) as u64,
729 );
730 }
731 }
732
733 fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx) {
734 let (n_in, n_out) = (a.n(), b[0].n());
735
736 debug_assert!(
737 n_out < n_in,
738 "invalid a: output ring degree should be smaller"
739 );
740 b[1..].iter().for_each(|bi| {
741 debug_assert_eq!(
742 bi.n(),
743 n_out,
744 "invalid input a: all VecZnx must have the same degree"
745 )
746 });
747
748 b.iter_mut().enumerate().for_each(|(i, bi)| {
749 if i == 0 {
750 switch_degree(bi, a);
751 self.vec_znx_rotate(-1, buf, a);
752 } else {
753 switch_degree(bi, buf);
754 self.vec_znx_rotate_inplace(-1, buf);
755 }
756 })
757 }
758
759 fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>) {
760 let (n_in, n_out) = (b.n(), a[0].n());
761
762 debug_assert!(
763 n_out < n_in,
764 "invalid a: output ring degree should be smaller"
765 );
766 a[1..].iter().for_each(|ai| {
767 debug_assert_eq!(
768 ai.n(),
769 n_out,
770 "invalid input a: all VecZnx must have the same degree"
771 )
772 });
773
774 a.iter().enumerate().for_each(|(_, ai)| {
775 switch_degree(b, ai);
776 self.vec_znx_rotate_inplace(-1, b);
777 });
778
779 self.vec_znx_rotate_inplace(a.len() as i64, b);
780 }
781}