1use crate::DataView;
  2use crate::DataViewMut;
  3use crate::ScalarZnx;
  4use crate::ZnxSliceSize;
  5use crate::ZnxZero;
  6use crate::alloc_aligned;
  7use crate::assert_alignement;
  8use crate::cast_mut;
  9use crate::ffi::znx;
 10use crate::znx_base::{ZnxInfos, ZnxView, ZnxViewMut};
 11use std::{cmp::min, fmt};
 12
 13/// [VecZnx] represents collection of contiguously stacked vector of small norm polynomials of
 14/// Zn\[X\] with [i64] coefficients.
 15/// A [VecZnx] is composed of multiple Zn\[X\] polynomials stored in a single contiguous array
 16/// in the memory.
 17///
 18/// # Example
 19///
 20/// Given 3 polynomials (a, b, c) of Zn\[X\], each with 4 columns, then the memory
 21/// layout is: `[a0, b0, c0, a1, b1, c1, a2, b2, c2, a3, b3, c3]`, where ai, bi, ci
 22/// are small polynomials of Zn\[X\].
 23pub struct VecZnx<D> {
 24    pub data: D,
 25    pub n: usize,
 26    pub cols: usize,
 27    pub size: usize,
 28}
 29
 30impl<D> ZnxInfos for VecZnx<D> {
 31    fn cols(&self) -> usize {
 32        self.cols
 33    }
 34
 35    fn rows(&self) -> usize {
 36        1
 37    }
 38
 39    fn n(&self) -> usize {
 40        self.n
 41    }
 42
 43    fn size(&self) -> usize {
 44        self.size
 45    }
 46}
 47
 48impl<D> ZnxSliceSize for VecZnx<D> {
 49    fn sl(&self) -> usize {
 50        self.n() * self.cols()
 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(&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    /// Truncates the precision of the [VecZnx] by k bits.
 73    ///
 74    /// # Arguments
 75    ///
 76    /// * `log_base2k`: the base two logarithm of the coefficients decomposition.
 77    /// * `k`: the number of bits of precision to drop.
 78    pub fn trunc_pow2(&mut self, log_base2k: usize, k: usize, col: usize) {
 79        if k == 0 {
 80            return;
 81        }
 82
 83        self.size -= k / log_base2k;
 84
 85        let k_rem: usize = k % log_base2k;
 86
 87        if k_rem != 0 {
 88            let mask: i64 = ((1 << (log_base2k - k_rem - 1)) - 1) << k_rem;
 89            self.at_mut(col, self.size() - 1)
 90                .iter_mut()
 91                .for_each(|x: &mut i64| *x &= mask)
 92        }
 93    }
 94}
 95
 96impl<D: From<Vec<u8>>> VecZnx<D> {
 97    pub(crate) fn bytes_of<Scalar: Sized>(n: usize, cols: usize, size: usize) -> usize {
 98        n * cols * size * size_of::<Scalar>()
 99    }
100
101    pub(crate) fn new<Scalar: Sized>(n: usize, cols: usize, size: usize) -> Self {
102        let data = alloc_aligned::<u8>(Self::bytes_of::<Scalar>(n, cols, size));
103        Self {
104            data: data.into(),
105            n,
106            cols,
107            size,
108        }
109    }
110
111    pub(crate) fn new_from_bytes<Scalar: Sized>(n: usize, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
112        let data: Vec<u8> = bytes.into();
113        assert!(data.len() == Self::bytes_of::<Scalar>(n, cols, size));
114        Self {
115            data: data.into(),
116            n,
117            cols,
118            size,
119        }
120    }
121}
122
123impl<D> VecZnx<D> {
124    pub(crate) fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
125        Self {
126            data,
127            n,
128            cols,
129            size,
130        }
131    }
132
133    pub fn to_scalar_znx(self) -> ScalarZnx<D> {
134        debug_assert_eq!(
135            self.size, 1,
136            "cannot convert VecZnx to ScalarZnx if cols: {} != 1",
137            self.cols
138        );
139        ScalarZnx {
140            data: self.data,
141            n: self.n,
142            cols: self.cols,
143        }
144    }
145}
146
147/// Copies the coefficients of `a` on the receiver.
148/// Copy is done with the minimum size matching both backing arrays.
149/// Panics if the cols do not match.
150pub fn copy_vec_znx_from<DataMut, Data>(b: &mut VecZnx<DataMut>, a: &VecZnx<Data>)
151where
152    DataMut: AsMut<[u8]> + AsRef<[u8]>,
153    Data: AsRef<[u8]>,
154{
155    assert_eq!(b.cols(), a.cols());
156    let data_a: &[i64] = a.raw();
157    let data_b: &mut [i64] = b.raw_mut();
158    let size = min(data_b.len(), data_a.len());
159    data_b[..size].copy_from_slice(&data_a[..size])
160}
161
162#[allow(dead_code)]
163fn normalize_tmp_bytes(n: usize) -> usize {
164    n * std::mem::size_of::<i64>()
165}
166
167#[allow(dead_code)]
168fn normalize<D: AsMut<[u8]> + AsRef<[u8]>>(log_base2k: usize, a: &mut VecZnx<D>, a_col: usize, tmp_bytes: &mut [u8]) {
169    let n: usize = a.n();
170
171    debug_assert!(
172        tmp_bytes.len() >= normalize_tmp_bytes(n),
173        "invalid tmp_bytes: tmp_bytes.len()={} < normalize_tmp_bytes({})",
174        tmp_bytes.len(),
175        n,
176    );
177
178    #[cfg(debug_assertions)]
179    {
180        assert_alignement(tmp_bytes.as_ptr())
181    }
182
183    let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
184
185    unsafe {
186        znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
187        (0..a.size()).rev().for_each(|i| {
188            znx::znx_normalize(
189                n as u64,
190                log_base2k as u64,
191                a.at_mut_ptr(a_col, i),
192                carry_i64.as_mut_ptr(),
193                a.at_mut_ptr(a_col, i),
194                carry_i64.as_mut_ptr(),
195            )
196        });
197    }
198}
199
200impl<D> VecZnx<D>
201where
202    VecZnx<D>: VecZnxToMut + ZnxInfos,
203{
204    /// Extracts the a_col-th column of 'a' and stores it on the self_col-th column [Self].
205    pub fn extract_column<R>(&mut self, self_col: usize, a: &R, a_col: usize)
206    where
207        R: VecZnxToRef + ZnxInfos,
208    {
209        #[cfg(debug_assertions)]
210        {
211            assert!(self_col < self.cols());
212            assert!(a_col < a.cols());
213        }
214
215        let min_size: usize = self.size.min(a.size());
216        let max_size: usize = self.size;
217
218        let mut self_mut: VecZnx<&mut [u8]> = self.to_mut();
219        let a_ref: VecZnx<&[u8]> = a.to_ref();
220
221        (0..min_size).for_each(|i: usize| {
222            self_mut
223                .at_mut(self_col, i)
224                .copy_from_slice(a_ref.at(a_col, i));
225        });
226
227        (min_size..max_size).for_each(|i| {
228            self_mut.zero_at(self_col, i);
229        });
230    }
231}
232
233impl<D: AsRef<[u8]>> fmt::Display for VecZnx<D> {
234    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
235        writeln!(
236            f,
237            "VecZnx(n={}, cols={}, size={})",
238            self.n, self.cols, self.size
239        )?;
240
241        for col in 0..self.cols {
242            writeln!(f, "Column {}:", col)?;
243            for size in 0..self.size {
244                let coeffs = self.at(col, size);
245                write!(f, "  Size {}: [", size)?;
246
247                let max_show = 100;
248                let show_count = coeffs.len().min(max_show);
249
250                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
251                    if i > 0 {
252                        write!(f, ", ")?;
253                    }
254                    write!(f, "{}", coeff)?;
255                }
256
257                if coeffs.len() > max_show {
258                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
259                }
260
261                writeln!(f, "]")?;
262            }
263        }
264        Ok(())
265    }
266}
267
268pub type VecZnxOwned = VecZnx<Vec<u8>>;
269pub type VecZnxMut<'a> = VecZnx<&'a mut [u8]>;
270pub type VecZnxRef<'a> = VecZnx<&'a [u8]>;
271
272pub trait VecZnxToRef {
273    fn to_ref(&self) -> VecZnx<&[u8]>;
274}
275
276pub trait VecZnxToMut {
277    fn to_mut(&mut self) -> VecZnx<&mut [u8]>;
278}
279
280impl VecZnxToMut for VecZnx<Vec<u8>> {
281    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
282        VecZnx {
283            data: self.data.as_mut_slice(),
284            n: self.n,
285            cols: self.cols,
286            size: self.size,
287        }
288    }
289}
290
291impl VecZnxToRef for VecZnx<Vec<u8>> {
292    fn to_ref(&self) -> VecZnx<&[u8]> {
293        VecZnx {
294            data: self.data.as_slice(),
295            n: self.n,
296            cols: self.cols,
297            size: self.size,
298        }
299    }
300}
301
302impl VecZnxToMut for VecZnx<&mut [u8]> {
303    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
304        VecZnx {
305            data: self.data,
306            n: self.n,
307            cols: self.cols,
308            size: self.size,
309        }
310    }
311}
312
313impl VecZnxToRef for VecZnx<&mut [u8]> {
314    fn to_ref(&self) -> VecZnx<&[u8]> {
315        VecZnx {
316            data: self.data,
317            n: self.n,
318            cols: self.cols,
319            size: self.size,
320        }
321    }
322}
323
324impl VecZnxToRef for VecZnx<&[u8]> {
325    fn to_ref(&self) -> VecZnx<&[u8]> {
326        VecZnx {
327            data: self.data,
328            n: self.n,
329            cols: self.cols,
330            size: self.size,
331        }
332    }
333}