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