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
136impl<D: From<Vec<u8>>> VecZnx<D> {
137    pub(crate) fn bytes_of<Scalar: Sized>(n: usize, cols: usize, size: usize) -> usize {
138        n * cols * size * size_of::<Scalar>()
139    }
140
141    pub(crate) fn new<Scalar: Sized>(n: usize, cols: usize, size: usize) -> Self {
142        let data = alloc_aligned::<u8>(Self::bytes_of::<Scalar>(n, cols, size));
143        Self {
144            data: data.into(),
145            n,
146            cols,
147            size,
148        }
149    }
150
151    pub(crate) fn new_from_bytes<Scalar: Sized>(n: usize, cols: usize, size: usize, bytes: impl Into<Vec<u8>>) -> Self {
152        let data: Vec<u8> = bytes.into();
153        assert!(data.len() == Self::bytes_of::<Scalar>(n, cols, size));
154        Self {
155            data: data.into(),
156            n,
157            cols,
158            size,
159        }
160    }
161}
162
163impl<D> VecZnx<D> {
164    pub(crate) fn from_data(data: D, n: usize, cols: usize, size: usize) -> Self {
165        Self {
166            data,
167            n,
168            cols,
169            size,
170        }
171    }
172
173    pub fn to_scalar_znx(self) -> ScalarZnx<D> {
174        debug_assert_eq!(
175            self.size, 1,
176            "cannot convert VecZnx to ScalarZnx if cols: {} != 1",
177            self.cols
178        );
179        ScalarZnx {
180            data: self.data,
181            n: self.n,
182            cols: self.cols,
183        }
184    }
185}
186
187/// Copies the coefficients of `a` on the receiver.
188/// Copy is done with the minimum size matching both backing arrays.
189/// Panics if the cols do not match.
190pub fn copy_vec_znx_from<DataMut, Data>(b: &mut VecZnx<DataMut>, a: &VecZnx<Data>)
191where
192    DataMut: AsMut<[u8]> + AsRef<[u8]>,
193    Data: AsRef<[u8]>,
194{
195    assert_eq!(b.cols(), a.cols());
196    let data_a: &[i64] = a.raw();
197    let data_b: &mut [i64] = b.raw_mut();
198    let size = min(data_b.len(), data_a.len());
199    data_b[..size].copy_from_slice(&data_a[..size])
200}
201
202#[allow(dead_code)]
203fn normalize_tmp_bytes(n: usize) -> usize {
204    n * std::mem::size_of::<i64>()
205}
206
207#[allow(dead_code)]
208fn normalize<D: AsMut<[u8]> + AsRef<[u8]>>(basek: usize, a: &mut VecZnx<D>, a_col: usize, tmp_bytes: &mut [u8]) {
209    let n: usize = a.n();
210
211    debug_assert!(
212        tmp_bytes.len() >= normalize_tmp_bytes(n),
213        "invalid tmp_bytes: tmp_bytes.len()={} < normalize_tmp_bytes({})",
214        tmp_bytes.len(),
215        n,
216    );
217
218    #[cfg(debug_assertions)]
219    {
220        assert_alignement(tmp_bytes.as_ptr())
221    }
222
223    let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
224
225    unsafe {
226        znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
227        (0..a.size()).rev().for_each(|i| {
228            znx::znx_normalize(
229                n as u64,
230                basek as u64,
231                a.at_mut_ptr(a_col, i),
232                carry_i64.as_mut_ptr(),
233                a.at_mut_ptr(a_col, i),
234                carry_i64.as_mut_ptr(),
235            )
236        });
237    }
238}
239
240impl<D> VecZnx<D>
241where
242    VecZnx<D>: VecZnxToMut + ZnxInfos,
243{
244    /// Extracts the a_col-th column of 'a' and stores it on the self_col-th column [Self].
245    pub fn extract_column<R>(&mut self, self_col: usize, a: &R, a_col: usize)
246    where
247        R: VecZnxToRef + ZnxInfos,
248    {
249        #[cfg(debug_assertions)]
250        {
251            assert!(self_col < self.cols());
252            assert!(a_col < a.cols());
253        }
254
255        let min_size: usize = self.size.min(a.size());
256        let max_size: usize = self.size;
257
258        let mut self_mut: VecZnx<&mut [u8]> = self.to_mut();
259        let a_ref: VecZnx<&[u8]> = a.to_ref();
260
261        (0..min_size).for_each(|i: usize| {
262            self_mut
263                .at_mut(self_col, i)
264                .copy_from_slice(a_ref.at(a_col, i));
265        });
266
267        (min_size..max_size).for_each(|i| {
268            self_mut.zero_at(self_col, i);
269        });
270    }
271}
272
273impl<D: AsRef<[u8]>> fmt::Display for VecZnx<D> {
274    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
275        writeln!(
276            f,
277            "VecZnx(n={}, cols={}, size={})",
278            self.n, self.cols, self.size
279        )?;
280
281        for col in 0..self.cols {
282            writeln!(f, "Column {}:", col)?;
283            for size in 0..self.size {
284                let coeffs = self.at(col, size);
285                write!(f, "  Size {}: [", size)?;
286
287                let max_show = 100;
288                let show_count = coeffs.len().min(max_show);
289
290                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
291                    if i > 0 {
292                        write!(f, ", ")?;
293                    }
294                    write!(f, "{}", coeff)?;
295                }
296
297                if coeffs.len() > max_show {
298                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
299                }
300
301                writeln!(f, "]")?;
302            }
303        }
304        Ok(())
305    }
306}
307
308pub type VecZnxOwned = VecZnx<Vec<u8>>;
309pub type VecZnxMut<'a> = VecZnx<&'a mut [u8]>;
310pub type VecZnxRef<'a> = VecZnx<&'a [u8]>;
311
312pub trait VecZnxToRef {
313    fn to_ref(&self) -> VecZnx<&[u8]>;
314}
315
316pub trait VecZnxToMut: VecZnxToRef {
317    fn to_mut(&mut self) -> VecZnx<&mut [u8]>;
318}
319
320impl VecZnxToMut for VecZnx<Vec<u8>> {
321    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
322        VecZnx {
323            data: self.data.as_mut_slice(),
324            n: self.n,
325            cols: self.cols,
326            size: self.size,
327        }
328    }
329}
330
331impl VecZnxToRef for VecZnx<Vec<u8>> {
332    fn to_ref(&self) -> VecZnx<&[u8]> {
333        VecZnx {
334            data: self.data.as_slice(),
335            n: self.n,
336            cols: self.cols,
337            size: self.size,
338        }
339    }
340}
341
342impl VecZnxToMut for VecZnx<&mut [u8]> {
343    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
344        VecZnx {
345            data: self.data,
346            n: self.n,
347            cols: self.cols,
348            size: self.size,
349        }
350    }
351}
352
353impl VecZnxToRef for VecZnx<&mut [u8]> {
354    fn to_ref(&self) -> VecZnx<&[u8]> {
355        VecZnx {
356            data: self.data,
357            n: self.n,
358            cols: self.cols,
359            size: self.size,
360        }
361    }
362}
363
364impl VecZnxToRef for VecZnx<&[u8]> {
365    fn to_ref(&self) -> VecZnx<&[u8]> {
366        VecZnx {
367            data: self.data,
368            n: self.n,
369            cols: self.cols,
370            size: self.size,
371        }
372    }
373}
374
375impl<DataSelf> VecZnx<DataSelf>
376where
377    VecZnx<DataSelf>: VecZnxToRef,
378{
379    pub fn clone(&self) -> VecZnx<Vec<u8>> {
380        let self_ref: VecZnx<&[u8]> = self.to_ref();
381
382        VecZnx {
383            data: self_ref.data.to_vec(),
384            n: self_ref.n,
385            cols: self_ref.cols,
386            size: self_ref.size,
387        }
388    }
389}