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 /// Allocates a new [VecZnx].
473 ///
474 /// # Arguments
475 ///
476 /// * `cols`: the number of cols.
477 fn new_vec_znx(&self, cols: usize) -> VecZnx;
478
479 /// Returns the minimum number of bytes necessary to allocate
480 /// a new [VecZnx] through [VecZnx::from_bytes].
481 fn bytes_of_vec_znx(&self, cols: usize) -> usize;
482
483 /// c <- a + b.
484 fn vec_znx_add<A: VecZnxCommon, B: VecZnxCommon, C: VecZnxCommon>(
485 &self,
486 c: &mut C,
487 a: &A,
488 b: &B,
489 );
490
491 /// b <- b + a.
492 fn vec_znx_add_inplace<A: VecZnxCommon, B: VecZnxCommon>(&self, b: &mut B, a: &A);
493
494 /// c <- a - b.
495 fn vec_znx_sub<A: VecZnxCommon, B: VecZnxCommon, C: VecZnxCommon>(
496 &self,
497 c: &mut C,
498 a: &A,
499 b: &B,
500 );
501
502 /// b <- b - a.
503 fn vec_znx_sub_inplace<A: VecZnxCommon, B: VecZnxCommon>(&self, b: &mut B, a: &A);
504
505 /// b <- -a.
506 fn vec_znx_negate<A: VecZnxCommon, B: VecZnxCommon>(&self, b: &mut B, a: &A);
507
508 /// b <- -b.
509 fn vec_znx_negate_inplace<A: VecZnxCommon>(&self, a: &mut A);
510
511 /// b <- a * X^k (mod X^{n} + 1)
512 fn vec_znx_rotate<A: VecZnxCommon, B: VecZnxCommon>(&self, k: i64, b: &mut B, a: &A);
513
514 /// a <- a * X^k (mod X^{n} + 1)
515 fn vec_znx_rotate_inplace<A: VecZnxCommon>(&self, k: i64, a: &mut A);
516
517 /// b <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
518 fn vec_znx_automorphism<A: VecZnxCommon, B: VecZnxCommon>(
519 &self,
520 k: i64,
521 b: &mut B,
522 a: &A,
523 a_cols: usize,
524 );
525
526 /// a <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
527 fn vec_znx_automorphism_inplace<A: VecZnxCommon>(&self, k: i64, a: &mut A, a_cols: usize);
528
529 /// Splits b into subrings and copies them them into a.
530 ///
531 /// # Panics
532 ///
533 /// This method requires that all [VecZnx] of b have the same ring degree
534 /// and that b.n() * b.len() <= a.n()
535 fn vec_znx_split<A: VecZnxCommon, B: VecZnxCommon, C: VecZnxCommon>(
536 &self,
537 b: &mut Vec<B>,
538 a: &A,
539 buf: &mut C,
540 );
541
542 /// Merges the subrings a into b.
543 ///
544 /// # Panics
545 ///
546 /// This method requires that all [VecZnx] of a have the same ring degree
547 /// and that a.n() * a.len() <= b.n()
548 fn vec_znx_merge<A: VecZnxCommon, B: VecZnxCommon>(&self, b: &mut B, a: &Vec<A>);
549}
550
551impl VecZnxOps for Module {
552 fn new_vec_znx(&self, cols: usize) -> VecZnx {
553 VecZnx::new(self.n(), cols)
554 }
555
556 fn bytes_of_vec_znx(&self, cols: usize) -> usize {
557 self.n() * cols * 8
558 }
559
560 // c <- a + b
561 fn vec_znx_add<A: VecZnxCommon, B: VecZnxCommon, C: VecZnxCommon>(
562 &self,
563 c: &mut C,
564 a: &A,
565 b: &B,
566 ) {
567 unsafe {
568 vec_znx::vec_znx_add(
569 self.0,
570 c.as_mut_ptr(),
571 c.cols() as u64,
572 c.n() as u64,
573 a.as_ptr(),
574 a.cols() as u64,
575 a.n() as u64,
576 b.as_ptr(),
577 b.cols() as u64,
578 b.n() as u64,
579 )
580 }
581 }
582
583 // b <- a + b
584 fn vec_znx_add_inplace<A: VecZnxCommon, B: VecZnxCommon>(&self, b: &mut B, a: &A) {
585 unsafe {
586 vec_znx::vec_znx_add(
587 self.0,
588 b.as_mut_ptr(),
589 b.cols() as u64,
590 b.n() as u64,
591 a.as_ptr(),
592 a.cols() as u64,
593 a.n() as u64,
594 b.as_ptr(),
595 b.cols() as u64,
596 b.n() as u64,
597 )
598 }
599 }
600
601 // c <- a + b
602 fn vec_znx_sub<A: VecZnxCommon, B: VecZnxCommon, C: VecZnxCommon>(
603 &self,
604 c: &mut C,
605 a: &A,
606 b: &B,
607 ) {
608 unsafe {
609 vec_znx::vec_znx_sub(
610 self.0,
611 c.as_mut_ptr(),
612 c.cols() as u64,
613 c.n() as u64,
614 a.as_ptr(),
615 a.cols() as u64,
616 a.n() as u64,
617 b.as_ptr(),
618 b.cols() as u64,
619 b.n() as u64,
620 )
621 }
622 }
623
624 // b <- a + b
625 fn vec_znx_sub_inplace<A: VecZnxCommon, B: VecZnxCommon>(&self, b: &mut B, a: &A) {
626 unsafe {
627 vec_znx::vec_znx_sub(
628 self.0,
629 b.as_mut_ptr(),
630 b.cols() as u64,
631 b.n() as u64,
632 a.as_ptr(),
633 a.cols() as u64,
634 a.n() as u64,
635 b.as_ptr(),
636 b.cols() as u64,
637 b.n() as u64,
638 )
639 }
640 }
641
642 fn vec_znx_negate<A: VecZnxCommon, B: VecZnxCommon>(&self, b: &mut B, a: &A) {
643 unsafe {
644 vec_znx::vec_znx_negate(
645 self.0,
646 b.as_mut_ptr(),
647 b.cols() as u64,
648 b.n() as u64,
649 a.as_ptr(),
650 a.cols() as u64,
651 a.n() as u64,
652 )
653 }
654 }
655
656 fn vec_znx_negate_inplace<A: VecZnxCommon>(&self, a: &mut A) {
657 unsafe {
658 vec_znx::vec_znx_negate(
659 self.0,
660 a.as_mut_ptr(),
661 a.cols() as u64,
662 a.n() as u64,
663 a.as_ptr(),
664 a.cols() as u64,
665 a.n() as u64,
666 )
667 }
668 }
669
670 fn vec_znx_rotate<A: VecZnxCommon, B: VecZnxCommon>(&self, k: i64, b: &mut B, a: &A) {
671 unsafe {
672 vec_znx::vec_znx_rotate(
673 self.0,
674 k,
675 b.as_mut_ptr(),
676 b.cols() as u64,
677 b.n() as u64,
678 a.as_ptr(),
679 a.cols() as u64,
680 a.n() as u64,
681 )
682 }
683 }
684
685 fn vec_znx_rotate_inplace<A: VecZnxCommon>(&self, k: i64, a: &mut A) {
686 unsafe {
687 vec_znx::vec_znx_rotate(
688 self.0,
689 k,
690 a.as_mut_ptr(),
691 a.cols() as u64,
692 a.n() as u64,
693 a.as_ptr(),
694 a.cols() as u64,
695 a.n() as u64,
696 )
697 }
698 }
699
700 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each cols.
701 ///
702 /// # Arguments
703 ///
704 /// * `a`: input.
705 /// * `b`: output.
706 /// * `k`: the power to which to map each coefficients.
707 /// * `a_cols`: the number of a_cols on which to apply the mapping.
708 ///
709 /// # Panics
710 ///
711 /// The method will panic if the argument `a` is greater than `a.cols()`.
712 ///
713 /// # Example
714 /// ```
715 /// use base2k::{Module, FFT64, VecZnx, Encoding, Infos, VecZnxApi, VecZnxOps};
716 /// use itertools::izip;
717 ///
718 /// let n: usize = 8; // polynomial degree
719 /// let module = Module::new::<FFT64>(n);
720 /// let mut a: VecZnx = module.new_vec_znx(2);
721 /// let mut b: VecZnx = module.new_vec_znx(2);
722 /// let mut c: VecZnx = module.new_vec_znx(2);
723 ///
724 /// (0..a.cols()).for_each(|i|{
725 /// a.at_mut(i).iter_mut().enumerate().for_each(|(i, x)|{
726 /// *x = i as i64
727 /// })
728 /// });
729 ///
730 /// module.vec_znx_automorphism(-1, &mut b, &a, 1); // X^i -> X^(-i)
731 /// let col = c.at_mut(0);
732 /// (1..col.len()).for_each(|i|{
733 /// col[n-i] = -(i as i64)
734 /// });
735 /// izip!(b.data.iter(), c.data.iter()).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
736 /// ```
737 fn vec_znx_automorphism<A: VecZnxCommon, B: VecZnxCommon>(
738 &self,
739 k: i64,
740 b: &mut B,
741 a: &A,
742 a_cols: usize,
743 ) {
744 assert_eq!(a.n(), self.n());
745 assert_eq!(b.n(), self.n());
746 assert!(a.cols() >= a_cols);
747 unsafe {
748 vec_znx::vec_znx_automorphism(
749 self.0,
750 k,
751 b.as_mut_ptr(),
752 b.cols() as u64,
753 b.n() as u64,
754 a.as_ptr(),
755 a_cols as u64,
756 a.n() as u64,
757 );
758 }
759 }
760
761 /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each cols.
762 ///
763 /// # Arguments
764 ///
765 /// * `a`: input and output.
766 /// * `k`: the power to which to map each coefficients.
767 /// * `a_cols`: the number of cols on which to apply the mapping.
768 ///
769 /// # Panics
770 ///
771 /// The method will panic if the argument `cols` is greater than `self.cols()`.
772 ///
773 /// # Example
774 /// ```
775 /// use base2k::{Module, FFT64, VecZnx, Encoding, Infos, VecZnxApi, VecZnxOps};
776 /// use itertools::izip;
777 ///
778 /// let n: usize = 8; // polynomial degree
779 /// let module = Module::new::<FFT64>(n);
780 /// let mut a: VecZnx = VecZnx::new(n, 2);
781 /// let mut b: VecZnx = VecZnx::new(n, 2);
782 ///
783 /// (0..a.cols()).for_each(|i|{
784 /// a.at_mut(i).iter_mut().enumerate().for_each(|(i, x)|{
785 /// *x = i as i64
786 /// })
787 /// });
788 ///
789 /// module.vec_znx_automorphism_inplace(-1, &mut a, 1); // X^i -> X^(-i)
790 /// let col = b.at_mut(0);
791 /// (1..col.len()).for_each(|i|{
792 /// col[n-i] = -(i as i64)
793 /// });
794 /// izip!(a.data.iter(), b.data.iter()).for_each(|(a, b)| assert_eq!(a, b, "{} != {}", a, b));
795 /// ```
796 fn vec_znx_automorphism_inplace<A: VecZnxCommon>(&self, k: i64, a: &mut A, a_cols: usize) {
797 assert_eq!(a.n(), self.n());
798 assert!(a.cols() >= a_cols);
799 unsafe {
800 vec_znx::vec_znx_automorphism(
801 self.0,
802 k,
803 a.as_mut_ptr(),
804 a.cols() as u64,
805 a.n() as u64,
806 a.as_ptr(),
807 a_cols as u64,
808 a.n() as u64,
809 );
810 }
811 }
812
813 fn vec_znx_split<A: VecZnxCommon, B: VecZnxCommon, C: VecZnxCommon>(
814 &self,
815 b: &mut Vec<B>,
816 a: &A,
817 buf: &mut C,
818 ) {
819 let (n_in, n_out) = (a.n(), b[0].n());
820
821 assert!(
822 n_out < n_in,
823 "invalid a: output ring degree should be smaller"
824 );
825 b[1..].iter().for_each(|bi| {
826 assert_eq!(
827 bi.n(),
828 n_out,
829 "invalid input a: all VecZnx must have the same degree"
830 )
831 });
832
833 b.iter_mut().enumerate().for_each(|(i, bi)| {
834 if i == 0 {
835 switch_degree(bi, a);
836 self.vec_znx_rotate(-1, buf, a);
837 } else {
838 switch_degree(bi, buf);
839 self.vec_znx_rotate_inplace(-1, buf);
840 }
841 })
842 }
843
844 fn vec_znx_merge<A: VecZnxCommon, B: VecZnxCommon>(&self, b: &mut B, a: &Vec<A>) {
845 let (n_in, n_out) = (b.n(), a[0].n());
846
847 assert!(
848 n_out < n_in,
849 "invalid a: output ring degree should be smaller"
850 );
851 a[1..].iter().for_each(|ai| {
852 assert_eq!(
853 ai.n(),
854 n_out,
855 "invalid input a: all VecZnx must have the same degree"
856 )
857 });
858
859 a.iter().enumerate().for_each(|(_, ai)| {
860 switch_degree(b, ai);
861 self.vec_znx_rotate_inplace(-1, b);
862 });
863
864 self.vec_znx_rotate_inplace(a.len() as i64, b);
865 }
866}