1use crate::Backend;
  2use crate::Module;
  3use crate::assert_alignement;
  4use crate::cast_mut;
  5use crate::ffi::znx;
  6use crate::znx_base::{GetZnxBase, ZnxAlloc, ZnxBase, ZnxBasics, ZnxInfos, ZnxLayout, ZnxSliceSize, switch_degree};
  7use std::cmp::min;
  8
  9pub const VEC_ZNX_ROWS: usize = 1;
 10
 11/// [VecZnx] represents collection of contiguously stacked vector of small norm polynomials of
 12/// Zn\[X\] with [i64] coefficients.
 13/// A [VecZnx] is composed of multiple Zn\[X\] polynomials stored in a single contiguous array
 14/// in the memory.
 15///
 16/// # Example
 17///
 18/// Given 3 polynomials (a, b, c) of Zn\[X\], each with 4 columns, then the memory
 19/// layout is: `[a0, b0, c0, a1, b1, c1, a2, b2, c2, a3, b3, c3]`, where ai, bi, ci
 20/// are small polynomials of Zn\[X\].
 21pub struct VecZnx {
 22    pub inner: ZnxBase,
 23}
 24
 25impl GetZnxBase for VecZnx {
 26    fn znx(&self) -> &ZnxBase {
 27        &self.inner
 28    }
 29
 30    fn znx_mut(&mut self) -> &mut ZnxBase {
 31        &mut self.inner
 32    }
 33}
 34
 35impl ZnxInfos for VecZnx {}
 36
 37impl ZnxSliceSize for VecZnx {
 38    fn sl(&self) -> usize {
 39        self.cols() * self.n()
 40    }
 41}
 42
 43impl ZnxLayout for VecZnx {
 44    type Scalar = i64;
 45}
 46
 47impl ZnxBasics for VecZnx {}
 48
 49impl<B: Backend> ZnxAlloc<B> for VecZnx {
 50    type Scalar = i64;
 51
 52    fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnx {
 53        debug_assert_eq!(bytes.len(), Self::bytes_of(module, _rows, cols, size));
 54        VecZnx {
 55            inner: ZnxBase::from_bytes_borrow(module.n(), VEC_ZNX_ROWS, cols, size, bytes),
 56        }
 57    }
 58
 59    fn bytes_of(module: &Module<B>, _rows: usize, cols: usize, size: usize) -> usize {
 60        debug_assert_eq!(
 61            _rows, VEC_ZNX_ROWS,
 62            "rows != {} not supported for VecZnx",
 63            VEC_ZNX_ROWS
 64        );
 65        module.n() * cols * size * size_of::<Self::Scalar>()
 66    }
 67}
 68
 69/// Copies the coefficients of `a` on the receiver.
 70/// Copy is done with the minimum size matching both backing arrays.
 71/// Panics if the cols do not match.
 72pub fn copy_vec_znx_from(b: &mut VecZnx, a: &VecZnx) {
 73    assert_eq!(b.cols(), a.cols());
 74    let data_a: &[i64] = a.raw();
 75    let data_b: &mut [i64] = b.raw_mut();
 76    let size = min(data_b.len(), data_a.len());
 77    data_b[..size].copy_from_slice(&data_a[..size])
 78}
 79
 80impl VecZnx {
 81    /// Truncates the precision of the [VecZnx] by k bits.
 82    ///
 83    /// # Arguments
 84    ///
 85    /// * `log_base2k`: the base two logarithm of the coefficients decomposition.
 86    /// * `k`: the number of bits of precision to drop.
 87    pub fn trunc_pow2(&mut self, log_base2k: usize, k: usize) {
 88        if k == 0 {
 89            return;
 90        }
 91
 92        if !self.borrowing() {
 93            self.inner
 94                .data
 95                .truncate(self.n() * self.cols() * (self.size() - k / log_base2k));
 96        }
 97
 98        self.inner.size -= k / log_base2k;
 99
100        let k_rem: usize = k % log_base2k;
101
102        if k_rem != 0 {
103            let mask: i64 = ((1 << (log_base2k - k_rem - 1)) - 1) << k_rem;
104            self.at_limb_mut(self.size() - 1)
105                .iter_mut()
106                .for_each(|x: &mut i64| *x &= mask)
107        }
108    }
109
110    pub fn copy_from(&mut self, a: &Self) {
111        copy_vec_znx_from(self, a);
112    }
113
114    pub fn normalize(&mut self, log_base2k: usize, carry: &mut [u8]) {
115        normalize(log_base2k, self, carry)
116    }
117
118    pub fn switch_degree(&self, col: usize, a: &mut Self, col_a: usize) {
119        switch_degree(a, col_a, self, col)
120    }
121
122    // Prints the first `n` coefficients of each limb
123    pub fn print(&self, n: usize) {
124        (0..self.size()).for_each(|i| println!("{}: {:?}", i, &self.at_limb(i)[..n]))
125    }
126}
127
128fn normalize_tmp_bytes(n: usize, size: usize) -> usize {
129    n * size * std::mem::size_of::<i64>()
130}
131
132fn normalize(log_base2k: usize, a: &mut VecZnx, tmp_bytes: &mut [u8]) {
133    let n: usize = a.n();
134    let cols: usize = a.cols();
135
136    debug_assert!(
137        tmp_bytes.len() >= normalize_tmp_bytes(n, cols),
138        "invalid tmp_bytes: tmp_bytes.len()={} < normalize_tmp_bytes({}, {})",
139        tmp_bytes.len(),
140        n,
141        cols,
142    );
143    #[cfg(debug_assertions)]
144    {
145        assert_alignement(tmp_bytes.as_ptr())
146    }
147
148    let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
149
150    unsafe {
151        znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
152        (0..a.size()).rev().for_each(|i| {
153            znx::znx_normalize(
154                (n * cols) as u64,
155                log_base2k as u64,
156                a.at_mut_ptr(0, i),
157                carry_i64.as_mut_ptr(),
158                a.at_mut_ptr(0, i),
159                carry_i64.as_mut_ptr(),
160            )
161        });
162    }
163}