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