1use crate::cast_mut;
  2use crate::ffi::vec_znx;
  3use crate::ffi::znx;
  4use crate::{alloc_aligned, assert_alignement};
  5use crate::{Infos, Module};
  6use itertools::izip;
  7use std::cmp::min;
  8
  9/// [VecZnx] represents a vector of small norm polynomials of Zn\[X\] with [i64] coefficients.
 10/// A [VecZnx] is composed of multiple Zn\[X\] polynomials stored in a single contiguous array
 11/// in the memory.
 12#[derive(Clone)]
 13pub struct VecZnx {
 14    /// Polynomial degree.
 15    pub n: usize,
 16
 17    /// Number of columns.
 18    pub cols: usize,
 19
 20    /// Polynomial coefficients, as a contiguous array. Each col is equally spaced by n.
 21    pub data: Vec<i64>,
 22
 23    /// Pointer to data (data can be enpty if [VecZnx] borrows space instead of owning it).
 24    pub ptr: *mut i64,
 25}
 26
 27pub trait VecZnxVec {
 28    fn dblptr(&self) -> Vec<&[i64]>;
 29    fn dblptr_mut(&mut self) -> Vec<&mut [i64]>;
 30}
 31
 32impl VecZnxVec for Vec<VecZnx> {
 33    fn dblptr(&self) -> Vec<&[i64]> {
 34        self.iter().map(|v| v.raw()).collect()
 35    }
 36
 37    fn dblptr_mut(&mut self) -> Vec<&mut [i64]> {
 38        self.iter_mut().map(|v| v.raw_mut()).collect()
 39    }
 40}
 41
 42pub fn bytes_of_vec_znx(n: usize, cols: usize) -> usize {
 43    n * cols * 8
 44}
 45
 46impl VecZnx {
 47    /// Returns a new struct implementing [VecZnx] with the provided data as backing array.
 48    ///
 49    /// The struct will take ownership of buf[..[VecZnx::bytes_of]]
 50    ///
 51    /// User must ensure that data is properly alligned and that
 52    /// the size of data is at least equal to [VecZnx::bytes_of].
 53    pub fn from_bytes(n: usize, cols: usize, bytes: &mut [u8]) -> Self {
 54        #[cfg(debug_assertions)]
 55        {
 56            assert_eq!(bytes.len(), Self::bytes_of(n, cols));
 57            assert_alignement(bytes.as_ptr());
 58        }
 59        unsafe {
 60            let bytes_i64: &mut [i64] = cast_mut::<u8, i64>(bytes);
 61            let ptr: *mut i64 = bytes_i64.as_mut_ptr();
 62            VecZnx {
 63                n: n,
 64                cols: cols,
 65                data: Vec::from_raw_parts(bytes_i64.as_mut_ptr(), bytes.len(), bytes.len()),
 66                ptr: ptr,
 67            }
 68        }
 69    }
 70
 71    pub fn from_bytes_borrow(n: usize, cols: usize, bytes: &mut [u8]) -> Self {
 72        #[cfg(debug_assertions)]
 73        {
 74            assert!(bytes.len() >= Self::bytes_of(n, cols));
 75            assert_alignement(bytes.as_ptr());
 76        }
 77        VecZnx {
 78            n: n,
 79            cols: cols,
 80            data: Vec::new(),
 81            ptr: bytes.as_mut_ptr() as *mut i64,
 82        }
 83    }
 84
 85    pub fn bytes_of(n: usize, cols: usize) -> usize {
 86        bytes_of_vec_znx(n, cols)
 87    }
 88
 89    pub fn copy_from(&mut self, a: &VecZnx) {
 90        copy_vec_znx_from(self, a);
 91    }
 92
 93    pub fn raw(&self) -> &[i64] {
 94        unsafe { std::slice::from_raw_parts(self.ptr, self.n * self.cols) }
 95    }
 96
 97    pub fn borrowing(&self) -> bool {
 98        self.data.len() == 0
 99    }
100
101    pub fn raw_mut(&mut self) -> &mut [i64] {
102        unsafe { std::slice::from_raw_parts_mut(self.ptr, self.n * self.cols) }
103    }
104
105    pub fn as_ptr(&self) -> *const i64 {
106        self.ptr
107    }
108
109    pub fn as_mut_ptr(&mut self) -> *mut i64 {
110        self.ptr
111    }
112
113    pub fn at(&self, i: usize) -> &[i64] {
114        let n: usize = self.n();
115        &self.raw()[n * i..n * (i + 1)]
116    }
117
118    pub fn at_mut(&mut self, i: usize) -> &mut [i64] {
119        let n: usize = self.n();
120        &mut self.raw_mut()[n * i..n * (i + 1)]
121    }
122
123    pub fn at_ptr(&self, i: usize) -> *const i64 {
124        self.ptr.wrapping_add(i * self.n)
125    }
126
127    pub fn at_mut_ptr(&mut self, i: usize) -> *mut i64 {
128        self.ptr.wrapping_add(i * self.n)
129    }
130
131    pub fn zero(&mut self) {
132        unsafe { znx::znx_zero_i64_ref((self.n * self.cols) as u64, self.ptr) }
133    }
134
135    pub fn normalize(&mut self, log_base2k: usize, carry: &mut [u8]) {
136        normalize(log_base2k, self, carry)
137    }
138
139    pub fn rsh(&mut self, log_base2k: usize, k: usize, carry: &mut [u8]) {
140        rsh(log_base2k, self, k, carry)
141    }
142
143    pub fn switch_degree(&self, a: &mut VecZnx) {
144        switch_degree(a, self)
145    }
146
147    pub fn print(&self, cols: usize, n: usize) {
148        (0..cols).for_each(|i| println!("{}: {:?}", i, &self.at(i)[..n]))
149    }
150}
151
152impl Infos for VecZnx {
153    /// Returns the base 2 logarithm of the [VecZnx] degree.
154    fn log_n(&self) -> usize {
155        (usize::BITS - (self.n - 1).leading_zeros()) as _
156    }
157
158    /// Returns the [VecZnx] degree.
159    fn n(&self) -> usize {
160        self.n
161    }
162
163    /// Returns the number of cols of the [VecZnx].
164    fn cols(&self) -> usize {
165        self.cols
166    }
167
168    /// Returns the number of rows of the [VecZnx].
169    fn rows(&self) -> usize {
170        1
171    }
172}
173
174/// Copies the coefficients of `a` on the receiver.
175/// Copy is done with the minimum size matching both backing arrays.
176pub fn copy_vec_znx_from(b: &mut VecZnx, a: &VecZnx) {
177    let data_a: &[i64] = a.raw();
178    let data_b: &mut [i64] = b.raw_mut();
179    let size = min(data_b.len(), data_a.len());
180    data_b[..size].copy_from_slice(&data_a[..size])
181}
182
183impl VecZnx {
184    /// Allocates a new [VecZnx] composed of #cols polynomials of Z\[X\].
185    pub fn new(n: usize, cols: usize) -> Self {
186        let mut data: Vec<i64> = alloc_aligned::<i64>(n * cols);
187        let ptr: *mut i64 = data.as_mut_ptr();
188        Self {
189            n: n,
190            cols: cols,
191            data: data,
192            ptr: ptr,
193        }
194    }
195
196    /// Truncates the precision of the [VecZnx] by k bits.
197    ///
198    /// # Arguments
199    ///
200    /// * `log_base2k`: the base two logarithm of the coefficients decomposition.
201    /// * `k`: the number of bits of precision to drop.
202    pub fn trunc_pow2(&mut self, log_base2k: usize, k: usize) {
203        if k == 0 {
204            return;
205        }
206
207        if !self.borrowing() {
208            self.data
209                .truncate((self.cols() - k / log_base2k) * self.n());
210        }
211
212        self.cols -= k / log_base2k;
213
214        let k_rem: usize = k % log_base2k;
215
216        if k_rem != 0 {
217            let mask: i64 = ((1 << (log_base2k - k_rem - 1)) - 1) << k_rem;
218            self.at_mut(self.cols() - 1)
219                .iter_mut()
220                .for_each(|x: &mut i64| *x &= mask)
221        }
222    }
223}
224
225pub fn switch_degree(b: &mut VecZnx, a: &VecZnx) {
226    let (n_in, n_out) = (a.n(), b.n());
227    let (gap_in, gap_out): (usize, usize);
228
229    if n_in > n_out {
230        (gap_in, gap_out) = (n_in / n_out, 1)
231    } else {
232        (gap_in, gap_out) = (1, n_out / n_in);
233        b.zero();
234    }
235
236    let cols = min(a.cols(), b.cols());
237
238    (0..cols).for_each(|i| {
239        izip!(
240            a.at(i).iter().step_by(gap_in),
241            b.at_mut(i).iter_mut().step_by(gap_out)
242        )
243        .for_each(|(x_in, x_out)| *x_out = *x_in);
244    });
245}
246
247fn normalize(log_base2k: usize, a: &mut VecZnx, tmp_bytes: &mut [u8]) {
248    let n: usize = a.n();
249
250    debug_assert!(
251        tmp_bytes.len() >= n * 8,
252        "invalid tmp_bytes: tmp_bytes.len()={} < self.n()={}",
253        tmp_bytes.len(),
254        n
255    );
256    #[cfg(debug_assertions)]
257    {
258        assert_alignement(tmp_bytes.as_ptr())
259    }
260
261    let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
262
263    unsafe {
264        znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
265        (0..a.cols()).rev().for_each(|i| {
266            znx::znx_normalize(
267                n as u64,
268                log_base2k as u64,
269                a.at_mut_ptr(i),
270                carry_i64.as_mut_ptr(),
271                a.at_mut_ptr(i),
272                carry_i64.as_mut_ptr(),
273            )
274        });
275    }
276}
277
278pub fn rsh(log_base2k: usize, a: &mut VecZnx, k: usize, tmp_bytes: &mut [u8]) {
279    let n: usize = a.n();
280
281    debug_assert!(
282        tmp_bytes.len() >> 3 >= n,
283        "invalid carry: carry.len()/8={} < self.n()={}",
284        tmp_bytes.len() >> 3,
285        n
286    );
287
288    #[cfg(debug_assertions)]
289    {
290        assert_alignement(tmp_bytes.as_ptr())
291    }
292
293    let cols: usize = a.cols();
294    let cols_steps: usize = k / log_base2k;
295
296    a.raw_mut().rotate_right(n * cols_steps);
297    unsafe {
298        znx::znx_zero_i64_ref((n * cols_steps) as u64, a.as_mut_ptr());
299    }
300
301    let k_rem = k % log_base2k;
302
303    if k_rem != 0 {
304        let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
305
306        unsafe {
307            znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
308        }
309
310        let log_base2k: usize = log_base2k;
311
312        (cols_steps..cols).for_each(|i| {
313            izip!(carry_i64.iter_mut(), a.at_mut(i).iter_mut()).for_each(|(ci, xi)| {
314                *xi += *ci << log_base2k;
315                *ci = get_base_k_carry(*xi, k_rem);
316                *xi = (*xi - *ci) >> k_rem;
317            });
318        })
319    }
320}
321
322#[inline(always)]
323fn get_base_k_carry(x: i64, k: usize) -> i64 {
324    (x << 64 - k) >> (64 - k)
325}
326
327pub trait VecZnxOps {
328    /// Allocates a new [VecZnx].
329    ///
330    /// # Arguments
331    ///
332    /// * `cols`: the number of cols.
333    fn new_vec_znx(&self, cols: usize) -> VecZnx;
334
335    /// Returns the minimum number of bytes necessary to allocate
336    /// a new [VecZnx] through [VecZnx::from_bytes].
337    fn bytes_of_vec_znx(&self, cols: usize) -> usize;
338
339    fn vec_znx_normalize_tmp_bytes(&self) -> usize;
340
341    /// c <- a + b.
342    fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
343
344    /// b <- b + a.
345    fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx);
346
347    /// c <- a - b.
348    fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx);
349
350    /// b <- b - a.
351    fn vec_znx_sub_inplace(&self, b: &mut VecZnx, a: &VecZnx);
352
353    /// b <- -a.
354    fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx);
355
356    /// b <- -b.
357    fn vec_znx_negate_inplace(&self, a: &mut VecZnx);
358
359    /// b <- a * X^k (mod X^{n} + 1)
360    fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
361
362    /// a <- a * X^k (mod X^{n} + 1)
363    fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx);
364
365    /// b <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
366    fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx);
367
368    /// a <- phi_k(a) where phi_k: X^i -> X^{i*k} (mod (X^{n} + 1))
369    fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx);
370
371    /// Splits b into subrings and copies them them into a.
372    ///
373    /// # Panics
374    ///
375    /// This method requires that all [VecZnx] of b have the same ring degree
376    /// and that b.n() * b.len() <= a.n()
377    fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx);
378
379    /// Merges the subrings a into b.
380    ///
381    /// # Panics
382    ///
383    /// This method requires that all [VecZnx] of a have the same ring degree
384    /// and that a.n() * a.len() <= b.n()
385    fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>);
386}
387
388impl VecZnxOps for Module {
389    fn new_vec_znx(&self, cols: usize) -> VecZnx {
390        VecZnx::new(self.n(), cols)
391    }
392
393    fn bytes_of_vec_znx(&self, cols: usize) -> usize {
394        self.n() * cols * 8
395    }
396
397    fn vec_znx_normalize_tmp_bytes(&self) -> usize {
398        unsafe { vec_znx::vec_znx_normalize_base2k_tmp_bytes(self.ptr) as usize }
399    }
400
401    // c <- a + b
402    fn vec_znx_add(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
403        unsafe {
404            vec_znx::vec_znx_add(
405                self.ptr,
406                c.as_mut_ptr(),
407                c.cols() as u64,
408                c.n() as u64,
409                a.as_ptr(),
410                a.cols() as u64,
411                a.n() as u64,
412                b.as_ptr(),
413                b.cols() as u64,
414                b.n() as u64,
415            )
416        }
417    }
418
419    // b <- a + b
420    fn vec_znx_add_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
421        unsafe {
422            vec_znx::vec_znx_add(
423                self.ptr,
424                b.as_mut_ptr(),
425                b.cols() as u64,
426                b.n() as u64,
427                a.as_ptr(),
428                a.cols() as u64,
429                a.n() as u64,
430                b.as_ptr(),
431                b.cols() as u64,
432                b.n() as u64,
433            )
434        }
435    }
436
437    // c <- a + b
438    fn vec_znx_sub(&self, c: &mut VecZnx, a: &VecZnx, b: &VecZnx) {
439        unsafe {
440            vec_znx::vec_znx_sub(
441                self.ptr,
442                c.as_mut_ptr(),
443                c.cols() as u64,
444                c.n() as u64,
445                a.as_ptr(),
446                a.cols() as u64,
447                a.n() as u64,
448                b.as_ptr(),
449                b.cols() as u64,
450                b.n() as u64,
451            )
452        }
453    }
454
455    // b <- a + b
456    fn vec_znx_sub_inplace(&self, b: &mut VecZnx, a: &VecZnx) {
457        unsafe {
458            vec_znx::vec_znx_sub(
459                self.ptr,
460                b.as_mut_ptr(),
461                b.cols() as u64,
462                b.n() as u64,
463                a.as_ptr(),
464                a.cols() as u64,
465                a.n() as u64,
466                b.as_ptr(),
467                b.cols() as u64,
468                b.n() as u64,
469            )
470        }
471    }
472
473    fn vec_znx_negate(&self, b: &mut VecZnx, a: &VecZnx) {
474        unsafe {
475            vec_znx::vec_znx_negate(
476                self.ptr,
477                b.as_mut_ptr(),
478                b.cols() as u64,
479                b.n() as u64,
480                a.as_ptr(),
481                a.cols() as u64,
482                a.n() as u64,
483            )
484        }
485    }
486
487    fn vec_znx_negate_inplace(&self, a: &mut VecZnx) {
488        unsafe {
489            vec_znx::vec_znx_negate(
490                self.ptr,
491                a.as_mut_ptr(),
492                a.cols() as u64,
493                a.n() as u64,
494                a.as_ptr(),
495                a.cols() as u64,
496                a.n() as u64,
497            )
498        }
499    }
500
501    fn vec_znx_rotate(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
502        unsafe {
503            vec_znx::vec_znx_rotate(
504                self.ptr,
505                k,
506                b.as_mut_ptr(),
507                b.cols() as u64,
508                b.n() as u64,
509                a.as_ptr(),
510                a.cols() as u64,
511                a.n() as u64,
512            )
513        }
514    }
515
516    fn vec_znx_rotate_inplace(&self, k: i64, a: &mut VecZnx) {
517        unsafe {
518            vec_znx::vec_znx_rotate(
519                self.ptr,
520                k,
521                a.as_mut_ptr(),
522                a.cols() as u64,
523                a.n() as u64,
524                a.as_ptr(),
525                a.cols() as u64,
526                a.n() as u64,
527            )
528        }
529    }
530
531    /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each cols.
532    ///
533    /// # Arguments
534    ///
535    /// * `a`: input.
536    /// * `b`: output.
537    /// * `k`: the power to which to map each coefficients.
538    /// * `a_cols`: the number of a_cols on which to apply the mapping.
539    ///
540    /// # Panics
541    ///
542    /// The method will panic if the argument `a` is greater than `a.cols()`.
543    fn vec_znx_automorphism(&self, k: i64, b: &mut VecZnx, a: &VecZnx) {
544        debug_assert_eq!(a.n(), self.n());
545        debug_assert_eq!(b.n(), self.n());
546        unsafe {
547            vec_znx::vec_znx_automorphism(
548                self.ptr,
549                k,
550                b.as_mut_ptr(),
551                b.cols() as u64,
552                b.n() as u64,
553                a.as_ptr(),
554                a.cols() as u64,
555                a.n() as u64,
556            );
557        }
558    }
559
560    /// Maps X^i to X^{ik} mod X^{n}+1. The mapping is applied independently on each cols.
561    ///
562    /// # Arguments
563    ///
564    /// * `a`: input and output.
565    /// * `k`: the power to which to map each coefficients.
566    /// * `a_cols`: the number of cols on which to apply the mapping.
567    ///
568    /// # Panics
569    ///
570    /// The method will panic if the argument `cols` is greater than `self.cols()`.
571    fn vec_znx_automorphism_inplace(&self, k: i64, a: &mut VecZnx) {
572        debug_assert_eq!(a.n(), self.n());
573        unsafe {
574            vec_znx::vec_znx_automorphism(
575                self.ptr,
576                k,
577                a.as_mut_ptr(),
578                a.cols() as u64,
579                a.n() as u64,
580                a.as_ptr(),
581                a.cols() as u64,
582                a.n() as u64,
583            );
584        }
585    }
586
587    fn vec_znx_split(&self, b: &mut Vec<VecZnx>, a: &VecZnx, buf: &mut VecZnx) {
588        let (n_in, n_out) = (a.n(), b[0].n());
589
590        debug_assert!(
591            n_out < n_in,
592            "invalid a: output ring degree should be smaller"
593        );
594        b[1..].iter().for_each(|bi| {
595            debug_assert_eq!(
596                bi.n(),
597                n_out,
598                "invalid input a: all VecZnx must have the same degree"
599            )
600        });
601
602        b.iter_mut().enumerate().for_each(|(i, bi)| {
603            if i == 0 {
604                switch_degree(bi, a);
605                self.vec_znx_rotate(-1, buf, a);
606            } else {
607                switch_degree(bi, buf);
608                self.vec_znx_rotate_inplace(-1, buf);
609            }
610        })
611    }
612
613    fn vec_znx_merge(&self, b: &mut VecZnx, a: &Vec<VecZnx>) {
614        let (n_in, n_out) = (b.n(), a[0].n());
615
616        debug_assert!(
617            n_out < n_in,
618            "invalid a: output ring degree should be smaller"
619        );
620        a[1..].iter().for_each(|ai| {
621            debug_assert_eq!(
622                ai.n(),
623                n_out,
624                "invalid input a: all VecZnx must have the same degree"
625            )
626        });
627
628        a.iter().enumerate().for_each(|(_, ai)| {
629            switch_degree(b, ai);
630            self.vec_znx_rotate_inplace(-1, b);
631        });
632
633        self.vec_znx_rotate_inplace(a.len() as i64, b);
634    }
635}