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    // Prints the first `n` coefficients of each limb
102    // pub fn print(&self, n: usize, col: usize) {
103    //     (0..self.size()).for_each(|j| println!("{}: {:?}", j, &self.at(col, j)[..n]));
104    // }
105}
106
107impl<D: From<Vec<u8>>> VecZnx<D> {
108    pub(crate) fn bytes_of<Scalar: Sized>(n: usize, cols: usize, size: usize) -> usize {
109        n * cols * size * size_of::<Scalar>()
110    }
111
112    pub(crate) fn new<Scalar: Sized>(n: usize, cols: usize, size: usize) -> Self {
113        let data = alloc_aligned::<u8>(Self::bytes_of::<Scalar>(n, cols, size));
114        Self {
115            data: data.into(),
116            n,
117            cols,
118            size,
119        }
120    }
121
122    pub(crate) fn new_from_bytes<Scalar: Sized>(n: usize, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
123        let data: Vec<u8> = bytes.into();
124        assert!(data.len() == Self::bytes_of::<Scalar>(n, cols, size));
125        Self {
126            data: data.into(),
127            n,
128            cols,
129            size,
130        }
131    }
132}
133
134//(Jay)TODO: Impl. truncate pow2 for Owned Vector
135
136/// Copies the coefficients of `a` on the receiver.
137/// Copy is done with the minimum size matching both backing arrays.
138/// Panics if the cols do not match.
139pub fn copy_vec_znx_from<DataMut, Data>(b: &mut VecZnx<DataMut>, a: &VecZnx<Data>)
140where
141    DataMut: AsMut<[u8]> + AsRef<[u8]>,
142    Data: AsRef<[u8]>,
143{
144    assert_eq!(b.cols(), a.cols());
145    let data_a: &[i64] = a.raw();
146    let data_b: &mut [i64] = b.raw_mut();
147    let size = min(data_b.len(), data_a.len());
148    data_b[..size].copy_from_slice(&data_a[..size])
149}
150
151// if !self.borrowing() {
152//     self.inner
153//         .data
154//         .truncate(self.n() * self.cols() * (self.size() - k / log_base2k));
155// }
156
157fn normalize_tmp_bytes(n: usize) -> usize {
158    n * std::mem::size_of::<i64>()
159}
160
161fn normalize<D: AsMut<[u8]> + AsRef<[u8]>>(log_base2k: usize, a: &mut VecZnx<D>, a_col: usize, tmp_bytes: &mut [u8]) {
162    let n: usize = a.n();
163
164    debug_assert!(
165        tmp_bytes.len() >= normalize_tmp_bytes(n),
166        "invalid tmp_bytes: tmp_bytes.len()={} < normalize_tmp_bytes({})",
167        tmp_bytes.len(),
168        n,
169    );
170
171    #[cfg(debug_assertions)]
172    {
173        assert_alignement(tmp_bytes.as_ptr())
174    }
175
176    let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
177
178    unsafe {
179        znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
180        (0..a.size()).rev().for_each(|i| {
181            znx::znx_normalize(
182                n as u64,
183                log_base2k as u64,
184                a.at_mut_ptr(a_col, i),
185                carry_i64.as_mut_ptr(),
186                a.at_mut_ptr(a_col, i),
187                carry_i64.as_mut_ptr(),
188            )
189        });
190    }
191}
192
193// impl<B: Backend> ZnxAlloc<B> for VecZnx {
194//     type Scalar = i64;
195
196//     fn from_bytes_borrow(module: &Module<B>, _rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> VecZnx {
197//         debug_assert_eq!(bytes.len(), Self::bytes_of(module, _rows, cols, size));
198//         VecZnx {
199//             inner: ZnxBase::from_bytes_borrow(module.n(), VEC_ZNX_ROWS, cols, size, bytes),
200//         }
201//     }
202
203//     fn bytes_of(module: &Module<B>, _rows: usize, cols: usize, size: usize) -> usize {
204//         debug_assert_eq!(
205//             _rows, VEC_ZNX_ROWS,
206//             "rows != {} not supported for VecZnx",
207//             VEC_ZNX_ROWS
208//         );
209//         module.n() * cols * size * size_of::<Self::Scalar>()
210//     }
211// }
212
213impl<D: AsRef<[u8]>> fmt::Display for VecZnx<D> {
214    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
215        writeln!(
216            f,
217            "VecZnx(n={}, cols={}, size={})",
218            self.n, self.cols, self.size
219        )?;
220
221        for col in 0..self.cols {
222            writeln!(f, "Column {}:", col)?;
223            for size in 0..self.size {
224                let coeffs = self.at(col, size);
225                write!(f, "  Size {}: [", size)?;
226
227                let max_show = 100;
228                let show_count = coeffs.len().min(max_show);
229
230                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
231                    if i > 0 {
232                        write!(f, ", ")?;
233                    }
234                    write!(f, "{}", coeff)?;
235                }
236
237                if coeffs.len() > max_show {
238                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
239                }
240
241                writeln!(f, "]")?;
242            }
243        }
244        Ok(())
245    }
246}
247
248pub type VecZnxOwned = VecZnx<Vec<u8>>;
249pub type VecZnxMut<'a> = VecZnx<&'a mut [u8]>;
250pub type VecZnxRef<'a> = VecZnx<&'a [u8]>;