1use crate::Backend;
  2use crate::ZnxBase;
  3use crate::cast_mut;
  4use crate::ffi::znx;
  5use crate::switch_degree;
  6use crate::{Module, ZnxBasics, ZnxInfos, ZnxLayout};
  7use crate::{alloc_aligned, assert_alignement};
  8use std::cmp::min;
  9
 10/// [VecZnx] represents collection of contiguously stacked vector of small norm polynomials of
 11/// Zn\[X\] with [i64] coefficients.
 12/// A [VecZnx] is composed of multiple Zn\[X\] polynomials stored in a single contiguous array
 13/// in the memory.
 14///
 15/// # Example
 16///
 17/// Given 3 polynomials (a, b, c) of Zn\[X\], each with 4 columns, then the memory
 18/// layout is: `[a0, b0, c0, a1, b1, c1, a2, b2, c2, a3, b3, c3]`, where ai, bi, ci
 19/// are small polynomials of Zn\[X\].
 20#[derive(Clone)]
 21pub struct VecZnx {
 22    /// Polynomial degree.
 23    pub n: usize,
 24
 25    /// The number of polynomials
 26    pub cols: usize,
 27
 28    /// The number of size per polynomial (a.k.a small polynomials).
 29    pub size: usize,
 30
 31    /// Polynomial coefficients, as a contiguous array. Each col is equally spaced by n.
 32    pub data: Vec<i64>,
 33
 34    /// Pointer to data (data can be enpty if [VecZnx] borrows space instead of owning it).
 35    pub ptr: *mut i64,
 36}
 37
 38impl ZnxInfos for VecZnx {
 39    fn n(&self) -> usize {
 40        self.n
 41    }
 42
 43    fn rows(&self) -> usize {
 44        1
 45    }
 46
 47    fn cols(&self) -> usize {
 48        self.cols
 49    }
 50
 51    fn size(&self) -> usize {
 52        self.size
 53    }
 54}
 55
 56impl ZnxLayout for VecZnx {
 57    type Scalar = i64;
 58
 59    fn as_ptr(&self) -> *const Self::Scalar {
 60        self.ptr
 61    }
 62
 63    fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
 64        self.ptr
 65    }
 66}
 67
 68impl ZnxBasics for VecZnx {}
 69
 70/// Copies the coefficients of `a` on the receiver.
 71/// Copy is done with the minimum size matching both backing arrays.
 72/// Panics if the cols do not match.
 73pub fn copy_vec_znx_from(b: &mut VecZnx, a: &VecZnx) {
 74    assert_eq!(b.cols(), a.cols());
 75    let data_a: &[i64] = a.raw();
 76    let data_b: &mut [i64] = b.raw_mut();
 77    let size = min(data_b.len(), data_a.len());
 78    data_b[..size].copy_from_slice(&data_a[..size])
 79}
 80
 81impl<B: Backend> ZnxBase<B> for VecZnx {
 82    type Scalar = i64;
 83
 84    /// Allocates a new [VecZnx] composed of #size polynomials of Z\[X\].
 85    fn new(module: &Module<B>, cols: usize, size: usize) -> Self {
 86        let n: usize = module.n();
 87        #[cfg(debug_assertions)]
 88        {
 89            assert!(n > 0);
 90            assert!(n & (n - 1) == 0);
 91            assert!(cols > 0);
 92            assert!(size > 0);
 93        }
 94        let mut data: Vec<i64> = alloc_aligned::<i64>(Self::bytes_of(module, cols, size));
 95        let ptr: *mut i64 = data.as_mut_ptr();
 96        Self {
 97            n: n,
 98            cols: cols,
 99            size: size,
100            data: data,
101            ptr: ptr,
102        }
103    }
104
105    fn bytes_of(module: &Module<B>, cols: usize, size: usize) -> usize {
106        module.n() * cols * size * size_of::<i64>()
107    }
108
109    /// Returns a new struct implementing [VecZnx] with the provided data as backing array.
110    ///
111    /// The struct will take ownership of buf[..[Self::bytes_of]]
112    ///
113    /// User must ensure that data is properly alligned and that
114    /// the size of data is equal to [Self::bytes_of].
115    fn from_bytes(module: &Module<B>, cols: usize, size: usize, bytes: &mut [u8]) -> Self {
116        let n: usize = module.n();
117        #[cfg(debug_assertions)]
118        {
119            assert!(cols > 0);
120            assert!(size > 0);
121            assert_eq!(bytes.len(), Self::bytes_of(module, cols, size));
122            assert_alignement(bytes.as_ptr());
123        }
124        unsafe {
125            let bytes_i64: &mut [i64] = cast_mut::<u8, i64>(bytes);
126            let ptr: *mut i64 = bytes_i64.as_mut_ptr();
127            Self {
128                n: n,
129                cols: cols,
130                size: size,
131                data: Vec::from_raw_parts(ptr, bytes.len(), bytes.len()),
132                ptr: ptr,
133            }
134        }
135    }
136
137    fn from_bytes_borrow(module: &Module<B>, cols: usize, size: usize, bytes: &mut [u8]) -> Self {
138        #[cfg(debug_assertions)]
139        {
140            assert!(cols > 0);
141            assert!(size > 0);
142            assert!(bytes.len() >= Self::bytes_of(module, cols, size));
143            assert_alignement(bytes.as_ptr());
144        }
145        Self {
146            n: module.n(),
147            cols: cols,
148            size: size,
149            data: Vec::new(),
150            ptr: bytes.as_mut_ptr() as *mut i64,
151        }
152    }
153}
154
155impl VecZnx {
156    /// Truncates the precision of the [VecZnx] by k bits.
157    ///
158    /// # Arguments
159    ///
160    /// * `log_base2k`: the base two logarithm of the coefficients decomposition.
161    /// * `k`: the number of bits of precision to drop.
162    pub fn trunc_pow2(&mut self, log_base2k: usize, k: usize) {
163        if k == 0 {
164            return;
165        }
166
167        if !self.borrowing() {
168            self.data
169                .truncate(self.n() * self.cols() * (self.size() - k / log_base2k));
170        }
171
172        self.size -= k / log_base2k;
173
174        let k_rem: usize = k % log_base2k;
175
176        if k_rem != 0 {
177            let mask: i64 = ((1 << (log_base2k - k_rem - 1)) - 1) << k_rem;
178            self.at_limb_mut(self.size() - 1)
179                .iter_mut()
180                .for_each(|x: &mut i64| *x &= mask)
181        }
182    }
183
184    pub fn copy_from(&mut self, a: &Self) {
185        copy_vec_znx_from(self, a);
186    }
187
188    pub fn borrowing(&self) -> bool {
189        self.data.len() == 0
190    }
191
192    pub fn normalize(&mut self, log_base2k: usize, carry: &mut [u8]) {
193        normalize(log_base2k, self, carry)
194    }
195
196    pub fn switch_degree(&self, a: &mut Self) {
197        switch_degree(a, self)
198    }
199
200    // Prints the first `n` coefficients of each limb
201    pub fn print(&self, n: usize) {
202        (0..self.size()).for_each(|i| println!("{}: {:?}", i, &self.at_limb(i)[..n]))
203    }
204}
205
206fn normalize_tmp_bytes(n: usize, size: usize) -> usize {
207    n * size * std::mem::size_of::<i64>()
208}
209
210fn normalize(log_base2k: usize, a: &mut VecZnx, tmp_bytes: &mut [u8]) {
211    let n: usize = a.n();
212    let cols: usize = a.cols();
213
214    debug_assert!(
215        tmp_bytes.len() >= normalize_tmp_bytes(n, cols),
216        "invalid tmp_bytes: tmp_bytes.len()={} < normalize_tmp_bytes({}, {})",
217        tmp_bytes.len(),
218        n,
219        cols,
220    );
221    #[cfg(debug_assertions)]
222    {
223        assert_alignement(tmp_bytes.as_ptr())
224    }
225
226    let carry_i64: &mut [i64] = cast_mut(tmp_bytes);
227
228    unsafe {
229        znx::znx_zero_i64_ref(n as u64, carry_i64.as_mut_ptr());
230        (0..a.size()).rev().for_each(|i| {
231            znx::znx_normalize(
232                (n * cols) as u64,
233                log_base2k as u64,
234                a.at_mut_ptr(0, i),
235                carry_i64.as_mut_ptr(),
236                a.at_mut_ptr(0, i),
237                carry_i64.as_mut_ptr(),
238            )
239        });
240    }
241}