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