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