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