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