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