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