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