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 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
246impl<D: AsRef<[u8]> + AsMut<[u8]>> VecZnx<D>{
247    pub fn normalize(&mut self, basek: usize, a_col: usize, tmp_bytes: &mut [u8]){
248        normalize(basek, self, a_col, tmp_bytes);
249    }
250}
251
252
253fn normalize<D: AsMut<[u8]> + AsRef<[u8]>>(basek: usize, a: &mut VecZnx<D>, a_col: usize, tmp_bytes: &mut [u8]) {
254    let n: usize = a.n();
255
256    debug_assert!(
257        tmp_bytes.len() >= normalize_tmp_bytes(n),
258        "invalid tmp_bytes: tmp_bytes.len()={} < normalize_tmp_bytes({})",
259        tmp_bytes.len(),
260        n,
261    );
262
263    #[cfg(debug_assertions)]
264    {
265        assert_alignement(tmp_bytes.as_ptr())
266    }
267
268    let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
269
270    unsafe {
271        znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
272        (0..a.size()).rev().for_each(|i| {
273            znx::znx_normalize(
274                n as u64,
275                basek as u64,
276                a.at_mut_ptr(a_col, i),
277                carry_i64.as_mut_ptr(),
278                a.at_mut_ptr(a_col, i),
279                carry_i64.as_mut_ptr(),
280            )
281        });
282    }
283}
284
285impl<D: AsMut<[u8]> + AsRef<[u8]>> VecZnx<D>
286where
287    VecZnx<D>: VecZnxToMut + ZnxInfos,
288{
289    /// Extracts the a_col-th column of 'a' and stores it on the self_col-th column [Self].
290    pub fn extract_column<R>(&mut self, self_col: usize, a: &VecZnx<R>, a_col: usize)
291    where
292        R: AsRef<[u8]>,
293        VecZnx<R>: VecZnxToRef + ZnxInfos,
294    {
295        #[cfg(debug_assertions)]
296        {
297            assert!(self_col < self.cols());
298            assert!(a_col < a.cols());
299        }
300
301        let min_size: usize = self.size.min(a.size());
302        let max_size: usize = self.size;
303
304        let mut self_mut: VecZnx<&mut [u8]> = self.to_mut();
305        let a_ref: VecZnx<&[u8]> = a.to_ref();
306
307        (0..min_size).for_each(|i: usize| {
308            self_mut
309                .at_mut(self_col, i)
310                .copy_from_slice(a_ref.at(a_col, i));
311        });
312
313        (min_size..max_size).for_each(|i| {
314            self_mut.zero_at(self_col, i);
315        });
316    }
317}
318
319impl<D: AsRef<[u8]>> fmt::Display for VecZnx<D> {
320    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
321        writeln!(
322            f,
323            "VecZnx(n={}, cols={}, size={})",
324            self.n, self.cols, self.size
325        )?;
326
327        for col in 0..self.cols {
328            writeln!(f, "Column {}:", col)?;
329            for size in 0..self.size {
330                let coeffs = self.at(col, size);
331                write!(f, "  Size {}: [", size)?;
332
333                let max_show = 100;
334                let show_count = coeffs.len().min(max_show);
335
336                for (i, &coeff) in coeffs.iter().take(show_count).enumerate() {
337                    if i > 0 {
338                        write!(f, ", ")?;
339                    }
340                    write!(f, "{}", coeff)?;
341                }
342
343                if coeffs.len() > max_show {
344                    write!(f, ", ... ({} more)", coeffs.len() - max_show)?;
345                }
346
347                writeln!(f, "]")?;
348            }
349        }
350        Ok(())
351    }
352}
353
354pub type VecZnxOwned = VecZnx<Vec<u8>>;
355pub type VecZnxMut<'a> = VecZnx<&'a mut [u8]>;
356pub type VecZnxRef<'a> = VecZnx<&'a [u8]>;
357
358pub trait VecZnxToRef {
359    fn to_ref(&self) -> VecZnx<&[u8]>;
360}
361
362impl<D> VecZnxToRef for VecZnx<D>
363where
364    D: AsRef<[u8]>,
365{
366    fn to_ref(&self) -> VecZnx<&[u8]> {
367        VecZnx {
368            data: self.data.as_ref(),
369            n: self.n,
370            cols: self.cols,
371            size: self.size,
372        }
373    }
374}
375
376pub trait VecZnxToMut {
377    fn to_mut(&mut self) -> VecZnx<&mut [u8]>;
378}
379
380impl<D> VecZnxToMut for VecZnx<D>
381where
382    D: AsRef<[u8]> + AsMut<[u8]>,
383{
384    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
385        VecZnx {
386            data: self.data.as_mut(),
387            n: self.n,
388            cols: self.cols,
389            size: self.size,
390        }
391    }
392}
393
394impl<DataSelf: AsRef<[u8]>> VecZnx<DataSelf> {
395    pub fn clone(&self) -> VecZnx<Vec<u8>> {
396        let self_ref: VecZnx<&[u8]> = self.to_ref();
397        VecZnx {
398            data: self_ref.data.to_vec(),
399            n: self_ref.n,
400            cols: self_ref.cols,
401            size: self_ref.size,
402        }
403    }
404}