1use crate::DataView;
  2use crate::DataViewMut;
  3use crate::alloc_aligned;
  4use crate::assert_alignement;
  5use crate::cast_mut;
  6use crate::ffi::znx;
  7use crate::znx_base::{ZnxInfos, ZnxView, ZnxViewMut, switch_degree};
  8use std::{cmp::min, fmt};
  9
 10// pub const VEC_ZNX_ROWS: usize = 1;
 11
 12/// [VecZnx] represents collection of contiguously stacked vector of small norm polynomials of
 13/// Zn\[X\] with [i64] coefficients.
 14/// A [VecZnx] is composed of multiple Zn\[X\] polynomials stored in a single contiguous array
 15/// in the memory.
 16///
 17/// # Example
 18///
 19/// Given 3 polynomials (a, b, c) of Zn\[X\], each with 4 columns, then the memory
 20/// layout is: `[a0, b0, c0, a1, b1, c1, a2, b2, c2, a3, b3, c3]`, where ai, bi, ci
 21/// are small polynomials of Zn\[X\].
 22pub struct VecZnx<D> {
 23    data: D,
 24    n: usize,
 25    cols: usize,
 26    size: usize,
 27}
 28
 29impl<D> ZnxInfos for VecZnx<D> {
 30    fn cols(&self) -> usize {
 31        self.cols
 32    }
 33
 34    fn rows(&self) -> usize {
 35        1
 36    }
 37
 38    fn n(&self) -> usize {
 39        self.n
 40    }
 41
 42    fn size(&self) -> usize {
 43        self.size
 44    }
 45
 46    fn sl(&self) -> usize {
 47        self.cols() * self.n()
 48    }
 49}
 50
 51impl<D> DataView for VecZnx<D> {
 52    type D = D;
 53    fn data(&self) -> &Self::D {
 54        &self.data
 55    }
 56}
 57
 58impl<D> DataViewMut for VecZnx<D> {
 59    fn data_mut(&mut self) -> &mut Self::D {
 60        &mut self.data
 61    }
 62}
 63
 64impl<D: AsRef<[u8]>> ZnxView for VecZnx<D> {
 65    type Scalar = i64;
 66}
 67
 68impl<D: AsMut<[u8]> + AsRef<[u8]>> VecZnx<D> {
 69    pub fn normalize(&mut self, log_base2k: usize, col: usize, carry: &mut [u8]) {
 70        normalize(log_base2k, self, col, carry)
 71    }
 72
 73    /// Truncates the precision of the [VecZnx] by k bits.
 74    ///
 75    /// # Arguments
 76    ///
 77    /// * `log_base2k`: the base two logarithm of the coefficients decomposition.
 78    /// * `k`: the number of bits of precision to drop.
 79    pub fn trunc_pow2(&mut self, log_base2k: usize, k: usize, col: usize) {
 80        if k == 0 {
 81            return;
 82        }
 83
 84        self.size -= k / log_base2k;
 85
 86        let k_rem: usize = k % log_base2k;
 87
 88        if k_rem != 0 {
 89            let mask: i64 = ((1 << (log_base2k - k_rem - 1)) - 1) << k_rem;
 90            self.at_mut(col, self.size() - 1)
 91                .iter_mut()
 92                .for_each(|x: &mut i64| *x &= mask)
 93        }
 94    }
 95
 96    /// Switches degree of from `a.n()` to `self.n()` into `self`
 97    pub fn switch_degree<Data: AsRef<[u8]>>(&mut self, col: usize, a: &VecZnx<Data>, col_a: usize) {
 98        switch_degree(self, col_a, a, col)
 99    }
100}
101
102impl<D: From<Vec<u8>>> VecZnx<D> {
103    pub(crate) fn bytes_of<Scalar: Sized>(n: usize, cols: usize, size: usize) -> usize {
104        n * cols * size * size_of::<Scalar>()
105    }
106
107    pub(crate) fn new<Scalar: Sized>(n: usize, cols: usize, size: usize) -> Self {
108        let data = alloc_aligned::<u8>(Self::bytes_of::<Scalar>(n, cols, size));
109        Self {
110            data: data.into(),
111            n,
112            cols,
113            size,
114        }
115    }
116
117    pub(crate) fn new_from_bytes<Scalar: Sized>(n: usize, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
118        let data: Vec<u8> = bytes.into();
119        assert!(data.len() == Self::bytes_of::<Scalar>(n, cols, size));
120        Self {
121            data: data.into(),
122            n,
123            cols,
124            size,
125        }
126    }
127}
128
129/// Copies the coefficients of `a` on the receiver.
130/// Copy is done with the minimum size matching both backing arrays.
131/// Panics if the cols do not match.
132pub fn copy_vec_znx_from<DataMut, Data>(b: &mut VecZnx<DataMut>, a: &VecZnx<Data>)
133where
134    DataMut: AsMut<[u8]> + AsRef<[u8]>,
135    Data: AsRef<[u8]>,
136{
137    assert_eq!(b.cols(), a.cols());
138    let data_a: &[i64] = a.raw();
139    let data_b: &mut [i64] = b.raw_mut();
140    let size = min(data_b.len(), data_a.len());
141    data_b[..size].copy_from_slice(&data_a[..size])
142}
143
144fn normalize_tmp_bytes(n: usize) -> usize {
145    n * std::mem::size_of::<i64>()
146}
147
148fn normalize<D: AsMut<[u8]> + AsRef<[u8]>>(log_base2k: usize, a: &mut VecZnx<D>, a_col: usize, tmp_bytes: &mut [u8]) {
149    let n: usize = a.n();
150
151    debug_assert!(
152        tmp_bytes.len() >= normalize_tmp_bytes(n),
153        "invalid tmp_bytes: tmp_bytes.len()={} < normalize_tmp_bytes({})",
154        tmp_bytes.len(),
155        n,
156    );
157
158    #[cfg(debug_assertions)]
159    {
160        assert_alignement(tmp_bytes.as_ptr())
161    }
162
163    let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
164
165    unsafe {
166        znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
167        (0..a.size()).rev().for_each(|i| {
168            znx::znx_normalize(
169                n as u64,
170                log_base2k as u64,
171                a.at_mut_ptr(a_col, i),
172                carry_i64.as_mut_ptr(),
173                a.at_mut_ptr(a_col, i),
174                carry_i64.as_mut_ptr(),
175            )
176        });
177    }
178}
179
180impl<D: AsRef<[u8]>> fmt::Display for VecZnx<D> {
181    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
182        writeln!(
183            f,
184            "VecZnx(n={}, cols={}, size={})",
185            self.n, self.cols, self.size
186        )?;
187
188        for col in 0..self.cols {
189            writeln!(f, "Column {}:", col)?;
190            for size in 0..self.size {
191                let coeffs = self.at(col, size);
192                write!(f, "  Size {}: [", size)?;
193
194                let max_show = 100;
195                let show_count = coeffs.len().min(max_show);
196
197                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
198                    if i > 0 {
199                        write!(f, ", ")?;
200                    }
201                    write!(f, "{}", coeff)?;
202                }
203
204                if coeffs.len() > max_show {
205                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
206                }
207
208                writeln!(f, "]")?;
209            }
210        }
211        Ok(())
212    }
213}
214
215pub type VecZnxOwned = VecZnx<Vec<u8>>;
216pub type VecZnxMut<'a> = VecZnx<&'a mut [u8]>;
217pub type VecZnxRef<'a> = VecZnx<&'a [u8]>;
218
219impl VecZnx<Vec<u8>> {
220    pub(crate) fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
221        VecZnx {
222            data: self.data.as_mut_slice(),
223            n: self.n,
224            cols: self.cols,
225            size: self.size,
226        }
227    }
228
229    pub(crate) fn to_ref(&self) -> VecZnx<&[u8]> {
230        VecZnx {
231            data: self.data.as_slice(),
232            n: self.n,
233            cols: self.cols,
234            size: self.size,
235        }
236    }
237}