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