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