1use crate::{Backend, Module, alloc_aligned, assert_alignement, cast_mut};
  2use itertools::izip;
  3use std::cmp::min;
  4
  5pub struct ZnxBase {
  6    /// The ring degree
  7    pub n: usize,
  8
  9    /// The number of rows (in the third dimension)
 10    pub rows: usize,
 11
 12    /// The number of polynomials
 13    pub cols: usize,
 14
 15    /// The number of size per polynomial (a.k.a small polynomials).
 16    pub size: usize,
 17
 18    /// Polynomial coefficients, as a contiguous array. Each col is equally spaced by n.
 19    pub data: Vec<u8>,
 20
 21    /// Pointer to data (data can be enpty if [VecZnx] borrows space instead of owning it).
 22    pub ptr: *mut u8,
 23}
 24
 25impl ZnxBase {
 26    pub fn from_bytes(n: usize, rows: usize, cols: usize, size: usize, mut bytes: Vec<u8>) -> Self {
 27        let mut res: Self = Self::from_bytes_borrow(n, rows, cols, size, &mut bytes);
 28        res.data = bytes;
 29        res
 30    }
 31
 32    pub fn from_bytes_borrow(n: usize, rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> Self {
 33        #[cfg(debug_assertions)]
 34        {
 35            assert_eq!(n & (n - 1), 0, "n must be a power of two");
 36            assert!(n > 0, "n must be greater than 0");
 37            assert!(rows > 0, "rows must be greater than 0");
 38            assert!(cols > 0, "cols must be greater than 0");
 39            assert!(size > 0, "size must be greater than 0");
 40        }
 41        Self {
 42            n: n,
 43            rows: rows,
 44            cols: cols,
 45            size: size,
 46            data: Vec::new(),
 47            ptr: bytes.as_mut_ptr(),
 48        }
 49    }
 50}
 51
 52pub trait GetZnxBase {
 53    fn znx(&self) -> &ZnxBase;
 54    fn znx_mut(&mut self) -> &mut ZnxBase;
 55}
 56
 57pub trait ZnxInfos {
 58    /// Returns the ring degree of the polynomials.
 59    fn n(&self) -> usize;
 60
 61    /// Returns the base two logarithm of the ring dimension of the polynomials.
 62    fn log_n(&self) -> usize {
 63        (usize::BITS - (self.n() - 1).leading_zeros()) as _
 64    }
 65
 66    /// Returns the number of rows.
 67    fn rows(&self) -> usize;
 68
 69    /// Returns the number of polynomials in each row.
 70    fn cols(&self) -> usize;
 71
 72    /// Returns the number of size per polynomial.
 73    fn size(&self) -> usize;
 74
 75    /// Returns the total number of small polynomials.
 76    fn poly_count(&self) -> usize {
 77        self.rows() * self.cols() * self.size()
 78    }
 79
 80    /// Returns the slice size, which is the offset between
 81    /// two size of the same column.
 82    fn sl(&self) -> usize;
 83}
 84
 85// pub trait ZnxSliceSize {}
 86
 87//(Jay) TODO: Remove ZnxAlloc
 88pub trait ZnxAlloc<B: Backend>
 89where
 90    Self: Sized + ZnxInfos,
 91{
 92    type Scalar;
 93    fn new(module: &Module<B>, rows: usize, cols: usize, size: usize) -> Self {
 94        let bytes: Vec<u8> = alloc_aligned::<u8>(Self::bytes_of(module, rows, cols, size));
 95        Self::from_bytes(module, rows, cols, size, bytes)
 96    }
 97
 98    fn from_bytes(module: &Module<B>, rows: usize, cols: usize, size: usize, mut bytes: Vec<u8>) -> Self {
 99        let mut res: Self = Self::from_bytes_borrow(module, rows, cols, size, &mut bytes);
100        res.znx_mut().data = bytes;
101        res
102    }
103
104    fn from_bytes_borrow(module: &Module<B>, rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> Self;
105
106    fn bytes_of(module: &Module<B>, rows: usize, cols: usize, size: usize) -> usize;
107}
108
109pub trait DataView {
110    type D;
111    fn data(&self) -> &Self::D;
112}
113
114pub trait DataViewMut: DataView {
115    fn data_mut(&self) -> &mut Self::D;
116}
117
118pub trait ZnxView: ZnxInfos + DataView<D: AsRef<[u8]>> {
119    type Scalar;
120
121    /// Returns a non-mutable pointer to the underlying coefficients array.
122    fn as_ptr(&self) -> *const Self::Scalar {
123        self.data().as_ref().as_ptr() as *const Self::Scalar
124    }
125
126    /// Returns a non-mutable reference to the entire underlying coefficient array.
127    fn raw(&self) -> &[Self::Scalar] {
128        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.n() * self.poly_count()) }
129    }
130
131    /// Returns a non-mutable pointer starting at the j-th small polynomial of the i-th column.
132    fn at_ptr(&self, i: usize, j: usize) -> *const Self::Scalar {
133        #[cfg(debug_assertions)]
134        {
135            assert!(i < self.cols());
136            assert!(j < self.size());
137        }
138        let offset: usize = self.n() * (j * self.cols() + i);
139        unsafe { self.as_ptr().add(offset) }
140    }
141
142    /// Returns non-mutable reference to the (i, j)-th small polynomial.
143    fn at(&self, i: usize, j: usize) -> &[Self::Scalar] {
144        unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n()) }
145    }
146}
147
148pub trait ZnxViewMut: ZnxView + DataViewMut<D: AsMut<[u8]>> {
149    /// Returns a mutable pointer to the underlying coefficients array.
150    fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
151        self.data_mut().as_mut().as_mut_ptr() as *mut Self::Scalar
152    }
153
154    /// Returns a mutable reference to the entire underlying coefficient array.
155    fn raw_mut(&mut self) -> &mut [Self::Scalar] {
156        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.n() * self.poly_count()) }
157    }
158
159    /// Returns a mutable pointer starting at the j-th small polynomial of the i-th column.
160    fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut Self::Scalar {
161        #[cfg(debug_assertions)]
162        {
163            assert!(i < self.cols());
164            assert!(j < self.size());
165        }
166        let offset: usize = self.n() * (j * self.cols() + i);
167        unsafe { self.as_mut_ptr().add(offset) }
168    }
169
170    /// Returns mutable reference to the (i, j)-th small polynomial.
171    fn at_mut(&mut self, i: usize, j: usize) -> &mut [Self::Scalar] {
172        unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n()) }
173    }
174}
175
176//(Jay)Note: Can't provide blanket impl. of ZnxView because Scalar is not known
177impl<T> ZnxViewMut for T where T: ZnxView + DataViewMut<D: AsMut<[u8]>> {}
178
179use std::convert::TryFrom;
180use std::num::TryFromIntError;
181use std::ops::{Add, AddAssign, Div, Mul, Neg, Shl, Shr, Sub};
182pub trait IntegerType:
183    Copy
184    + std::fmt::Debug
185    + Default
186    + PartialEq
187    + PartialOrd
188    + Add<Output = Self>
189    + Sub<Output = Self>
190    + Mul<Output = Self>
191    + Div<Output = Self>
192    + Neg<Output = Self>
193    + Shr<Output = Self>
194    + Shl<Output = Self>
195    + AddAssign
196    + TryFrom<usize, Error = TryFromIntError>
197{
198    const BITS: u32;
199}
200
201impl IntegerType for i64 {
202    const BITS: u32 = 64;
203}
204
205impl IntegerType for i128 {
206    const BITS: u32 = 128;
207}
208
209pub trait ZnxZero: ZnxViewMut
210where
211    Self: Sized,
212{
213    fn zero(&mut self) {
214        unsafe {
215            std::ptr::write_bytes(
216                self.as_mut_ptr(),
217                0,
218                self.n() * size_of::<Self::Scalar>() * self.poly_count(),
219            );
220        }
221    }
222
223    fn zero_at(&mut self, i: usize, j: usize) {
224        unsafe {
225            std::ptr::write_bytes(
226                self.at_mut_ptr(i, j),
227                0,
228                self.n() * size_of::<Self::Scalar>(),
229            );
230        }
231    }
232}
233
234pub trait ZnxRsh: ZnxZero {
235    fn rsh(&mut self, k: usize, log_base2k: usize, col: usize, carry: &mut [u8]) {
236        rsh(k, log_base2k, self, col, carry)
237    }
238}
239
240// Blanket implementations
241impl<T> ZnxZero for T where T: ZnxViewMut {}
242impl<T> ZnxRsh for T where T: ZnxZero {}
243
244pub fn rsh<V: ZnxRsh + ZnxZero>(k: usize, log_base2k: usize, a: &mut V, a_col: usize, tmp_bytes: &mut [u8])
245where
246    V::Scalar: IntegerType,
247{
248    let n: usize = a.n();
249    let size: usize = a.size();
250    let cols: usize = a.cols();
251
252    #[cfg(debug_assertions)]
253    {
254        assert!(
255            tmp_bytes.len() >= rsh_tmp_bytes::<V::Scalar>(n),
256            "invalid carry: carry.len()/size_ofSelf::Scalar={} < rsh_tmp_bytes({}, {})",
257            tmp_bytes.len() / size_of::<V::Scalar>(),
258            n,
259            size,
260        );
261        assert_alignement(tmp_bytes.as_ptr());
262    }
263
264    let size: usize = a.size();
265    let steps: usize = k / log_base2k;
266
267    a.raw_mut().rotate_right(n * steps * cols);
268    (0..cols).for_each(|i| {
269        (0..steps).for_each(|j| {
270            a.zero_at(i, j);
271        })
272    });
273
274    let k_rem: usize = k % log_base2k;
275
276    if k_rem != 0 {
277        let carry: &mut [V::Scalar] = cast_mut(tmp_bytes);
278
279        unsafe {
280            std::ptr::write_bytes(carry.as_mut_ptr(), 0, n * size_of::<V::Scalar>());
281        }
282
283        let log_base2k_t: V::Scalar = V::Scalar::try_from(log_base2k).unwrap();
284        let shift: V::Scalar = V::Scalar::try_from(V::Scalar::BITS as usize - k_rem).unwrap();
285        let k_rem_t: V::Scalar = V::Scalar::try_from(k_rem).unwrap();
286
287        (steps..size).for_each(|i| {
288            izip!(carry.iter_mut(), a.at_mut(a_col, i).iter_mut()).for_each(|(ci, xi)| {
289                *xi += *ci << log_base2k_t;
290                *ci = get_base_k_carry(*xi, shift);
291                *xi = (*xi - *ci) >> k_rem_t;
292            });
293        })
294    }
295}
296
297#[inline(always)]
298fn get_base_k_carry<T: IntegerType>(x: T, shift: T) -> T {
299    (x << shift) >> shift
300}
301
302pub fn rsh_tmp_bytes<T: IntegerType>(n: usize) -> usize {
303    n * std::mem::size_of::<T>()
304}
305
306pub fn switch_degree<DMut: ZnxViewMut + ZnxZero, D: ZnxView>(b: &mut DMut, col_b: usize, a: &D, col_a: usize) {
307    let (n_in, n_out) = (a.n(), b.n());
308    let (gap_in, gap_out): (usize, usize);
309
310    if n_in > n_out {
311        (gap_in, gap_out) = (n_in / n_out, 1)
312    } else {
313        (gap_in, gap_out) = (1, n_out / n_in);
314        b.zero();
315    }
316
317    let size: usize = min(a.size(), b.size());
318
319    (0..size).for_each(|i| {
320        izip!(
321            a.at(col_a, i).iter().step_by(gap_in),
322            b.at_mut(col_b, i).iter_mut().step_by(gap_out)
323        )
324        .for_each(|(x_in, x_out)| *x_out = *x_in);
325    });
326}
327
328// pub trait ZnxLayout: ZnxInfos {
329//     type Scalar;
330
331//     /// Returns true if the receiver is only borrowing the data.
332//     fn borrowing(&self) -> bool {
333//         self.znx().data.len() == 0
334//     }
335
336//     /// Returns a non-mutable pointer to the underlying coefficients array.
337//     fn as_ptr(&self) -> *const Self::Scalar {
338//         self.znx().ptr as *const Self::Scalar
339//     }
340
341//     /// Returns a mutable pointer to the underlying coefficients array.
342//     fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
343//         self.znx_mut().ptr as *mut Self::Scalar
344//     }
345
346//     /// Returns a non-mutable reference to the entire underlying coefficient array.
347//     fn raw(&self) -> &[Self::Scalar] {
348//         unsafe { std::slice::from_raw_parts(self.as_ptr(), self.n() * self.poly_count()) }
349//     }
350
351//     /// Returns a mutable reference to the entire underlying coefficient array.
352//     fn raw_mut(&mut self) -> &mut [Self::Scalar] {
353//         unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.n() * self.poly_count()) }
354//     }
355
356//     /// Returns a non-mutable pointer starting at the j-th small polynomial of the i-th column.
357//     fn at_ptr(&self, i: usize, j: usize) -> *const Self::Scalar {
358//         #[cfg(debug_assertions)]
359//         {
360//             assert!(i < self.cols());
361//             assert!(j < self.size());
362//         }
363//         let offset: usize = self.n() * (j * self.cols() + i);
364//         unsafe { self.as_ptr().add(offset) }
365//     }
366
367//     /// Returns a mutable pointer starting at the j-th small polynomial of the i-th column.
368//     fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut Self::Scalar {
369//         #[cfg(debug_assertions)]
370//         {
371//             assert!(i < self.cols());
372//             assert!(j < self.size());
373//         }
374//         let offset: usize = self.n() * (j * self.cols() + i);
375//         unsafe { self.as_mut_ptr().add(offset) }
376//     }
377
378//     /// Returns non-mutable reference to the (i, j)-th small polynomial.
379//     fn at(&self, i: usize, j: usize) -> &[Self::Scalar] {
380//         unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n()) }
381//     }
382
383//     /// Returns mutable reference to the (i, j)-th small polynomial.
384//     fn at_mut(&mut self, i: usize, j: usize) -> &mut [Self::Scalar] {
385//         unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n()) }
386//     }
387// }