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