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, ZnxInfos, ZnxLayout, ZnxRsh, ZnxSliceSize, ZnxZero, 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 ZnxZero for VecZnx {}
 48
 49impl ZnxRsh for VecZnx {}
 50
 51impl<B: Backend> ZnxAlloc<B> for VecZnx {
 52    type Scalar = i64;
 53
 54    fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnx {
 55        debug_assert_eq!(bytes.len(), Self::bytes_of(module, _rows, cols, size));
 56        VecZnx {
 57            inner: ZnxBase::from_bytes_borrow(module.n(), VEC_ZNX_ROWS, cols, size, bytes),
 58        }
 59    }
 60
 61    fn bytes_of(module: &Module<B>, _rows: usize, cols: usize, size: usize) -> usize {
 62        debug_assert_eq!(
 63            _rows, VEC_ZNX_ROWS,
 64            "rows != {} not supported for VecZnx",
 65            VEC_ZNX_ROWS
 66        );
 67        module.n() * cols * size * size_of::<Self::Scalar>()
 68    }
 69}
 70
 71/// Copies the coefficients of `a` on the receiver.
 72/// Copy is done with the minimum size matching both backing arrays.
 73/// Panics if the cols do not match.
 74pub fn copy_vec_znx_from(b: &mut VecZnx, a: &VecZnx) {
 75    assert_eq!(b.cols(), a.cols());
 76    let data_a: &[i64] = a.raw();
 77    let data_b: &mut [i64] = b.raw_mut();
 78    let size = min(data_b.len(), data_a.len());
 79    data_b[..size].copy_from_slice(&data_a[..size])
 80}
 81
 82impl VecZnx {
 83    /// Truncates the precision of the [VecZnx] by k bits.
 84    ///
 85    /// # Arguments
 86    ///
 87    /// * `log_base2k`: the base two logarithm of the coefficients decomposition.
 88    /// * `k`: the number of bits of precision to drop.
 89    pub fn trunc_pow2(&mut self, log_base2k: usize, k: usize, col: usize) {
 90        if k == 0 {
 91            return;
 92        }
 93
 94        if !self.borrowing() {
 95            self.inner
 96                .data
 97                .truncate(self.n() * self.cols() * (self.size() - k / log_base2k));
 98        }
 99
100        self.inner.size -= k / log_base2k;
101
102        let k_rem: usize = k % log_base2k;
103
104        if k_rem != 0 {
105            let mask: i64 = ((1 << (log_base2k - k_rem - 1)) - 1) << k_rem;
106            self.at_mut(col, self.size() - 1)
107                .iter_mut()
108                .for_each(|x: &mut i64| *x &= mask)
109        }
110    }
111
112    pub fn copy_from(&mut self, a: &Self) {
113        copy_vec_znx_from(self, a);
114    }
115
116    pub fn normalize(&mut self, log_base2k: usize, col: usize, carry: &mut [u8]) {
117        normalize(log_base2k, self, col, carry)
118    }
119
120    pub fn switch_degree(&self, col: usize, a: &mut Self, col_a: usize) {
121        switch_degree(a, col_a, self, col)
122    }
123
124    // Prints the first `n` coefficients of each limb
125    pub fn print(&self, n: usize, col: usize) {
126        (0..self.size()).for_each(|j| println!("{}: {:?}", j, &self.at(col, j)[..n]));
127    }
128}
129
130fn normalize_tmp_bytes(n: usize) -> usize {
131    n * std::mem::size_of::<i64>()
132}
133
134fn normalize(log_base2k: usize, a: &mut VecZnx, a_col: usize, tmp_bytes: &mut [u8]) {
135    let n: usize = a.n();
136
137    debug_assert!(
138        tmp_bytes.len() >= normalize_tmp_bytes(n),
139        "invalid tmp_bytes: tmp_bytes.len()={} < normalize_tmp_bytes({})",
140        tmp_bytes.len(),
141        n,
142    );
143
144    #[cfg(debug_assertions)]
145    {
146        assert_alignement(tmp_bytes.as_ptr())
147    }
148
149    let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
150
151    unsafe {
152        znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
153        (0..a.size()).rev().for_each(|i| {
154            znx::znx_normalize(
155                n as u64,
156                log_base2k as u64,
157                a.at_mut_ptr(a_col, i),
158                carry_i64.as_mut_ptr(),
159                a.at_mut_ptr(a_col, i),
160                carry_i64.as_mut_ptr(),
161            )
162        });
163    }
164}