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: GetZnxBase {
 58    /// Returns the ring degree of the polynomials.
 59    fn n(&self) -> usize {
 60        self.znx().n
 61    }
 62
 63    /// Returns the base two logarithm of the ring dimension of the polynomials.
 64    fn log_n(&self) -> usize {
 65        (usize::BITS - (self.n() - 1).leading_zeros()) as _
 66    }
 67
 68    /// Returns the number of rows.
 69    fn rows(&self) -> usize {
 70        self.znx().rows
 71    }
 72    /// Returns the number of polynomials in each row.
 73    fn cols(&self) -> usize {
 74        self.znx().cols
 75    }
 76
 77    /// Returns the number of size per polynomial.
 78    fn size(&self) -> usize {
 79        self.znx().size
 80    }
 81
 82    /// Returns the underlying raw bytes array.
 83    fn data(&self) -> &[u8] {
 84        &self.znx().data
 85    }
 86
 87    /// Returns a pointer to the underlying raw bytes array.
 88    fn ptr(&self) -> *mut u8 {
 89        self.znx().ptr
 90    }
 91
 92    /// Returns the total number of small polynomials.
 93    fn poly_count(&self) -> usize {
 94        self.rows() * self.cols() * self.size()
 95    }
 96}
 97
 98pub trait ZnxSliceSize {
 99    /// Returns the slice size, which is the offset between
100    /// two size of the same column.
101    fn sl(&self) -> usize;
102}
103
104pub trait ZnxAlloc<B: Backend>
105where
106    Self: Sized + ZnxInfos,
107{
108    type Scalar;
109    fn new(module: &Module<B>, rows: usize, cols: usize, size: usize) -> Self {
110        let bytes: Vec<u8> = alloc_aligned::<u8>(Self::bytes_of(module, rows, cols, size));
111        Self::from_bytes(module, rows, cols, size, bytes)
112    }
113
114    fn from_bytes(module: &Module<B>, rows: usize, cols: usize, size: usize, mut bytes: Vec<u8>) -> Self {
115        let mut res: Self = Self::from_bytes_borrow(module, rows, cols, size, &mut bytes);
116        res.znx_mut().data = bytes;
117        res
118    }
119
120    fn from_bytes_borrow(module: &Module<B>, rows: usize, cols: usize, size: usize, bytes: &mut [u8]) -> Self;
121
122    fn bytes_of(module: &Module<B>, rows: usize, cols: usize, size: usize) -> usize;
123}
124
125pub trait ZnxLayout: ZnxInfos {
126    type Scalar;
127
128    /// Returns true if the receiver is only borrowing the data.
129    fn borrowing(&self) -> bool {
130        self.znx().data.len() == 0
131    }
132
133    /// Returns a non-mutable pointer to the underlying coefficients array.
134    fn as_ptr(&self) -> *const Self::Scalar {
135        self.znx().ptr as *const Self::Scalar
136    }
137
138    /// Returns a mutable pointer to the underlying coefficients array.
139    fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
140        self.znx_mut().ptr as *mut Self::Scalar
141    }
142
143    /// Returns a non-mutable reference to the entire underlying coefficient array.
144    fn raw(&self) -> &[Self::Scalar] {
145        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.n() * self.poly_count()) }
146    }
147
148    /// Returns a mutable reference to the entire underlying coefficient array.
149    fn raw_mut(&mut self) -> &mut [Self::Scalar] {
150        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.n() * self.poly_count()) }
151    }
152
153    /// Returns a non-mutable pointer starting at the j-th small polynomial of the i-th column.
154    fn at_ptr(&self, i: usize, j: usize) -> *const Self::Scalar {
155        #[cfg(debug_assertions)]
156        {
157            assert!(i < self.cols());
158            assert!(j < self.size());
159        }
160        let offset: usize = self.n() * (j * self.cols() + i);
161        unsafe { self.as_ptr().add(offset) }
162    }
163
164    /// Returns a mutable pointer starting at the j-th small polynomial of the i-th column.
165    fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut Self::Scalar {
166        #[cfg(debug_assertions)]
167        {
168            assert!(i < self.cols());
169            assert!(j < self.size());
170        }
171        let offset: usize = self.n() * (j * self.cols() + i);
172        unsafe { self.as_mut_ptr().add(offset) }
173    }
174
175    /// Returns non-mutable reference to the (i, j)-th small polynomial.
176    fn at(&self, i: usize, j: usize) -> &[Self::Scalar] {
177        unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n()) }
178    }
179
180    /// Returns mutable reference to the (i, j)-th small polynomial.
181    fn at_mut(&mut self, i: usize, j: usize) -> &mut [Self::Scalar] {
182        unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n()) }
183    }
184}
185
186use std::convert::TryFrom;
187use std::num::TryFromIntError;
188use std::ops::{Add, AddAssign, Div, Mul, Neg, Shl, Shr, Sub};
189pub trait IntegerType:
190    Copy
191    + std::fmt::Debug
192    + Default
193    + PartialEq
194    + PartialOrd
195    + Add<Output = Self>
196    + Sub<Output = Self>
197    + Mul<Output = Self>
198    + Div<Output = Self>
199    + Neg<Output = Self>
200    + Shr<Output = Self>
201    + Shl<Output = Self>
202    + AddAssign
203    + TryFrom<usize, Error = TryFromIntError>
204{
205    const BITS: u32;
206}
207
208impl IntegerType for i64 {
209    const BITS: u32 = 64;
210}
211
212impl IntegerType for i128 {
213    const BITS: u32 = 128;
214}
215
216pub trait ZnxZero: ZnxLayout
217where
218    Self: Sized,
219{
220    fn zero(&mut self) {
221        unsafe {
222            std::ptr::write_bytes(
223                self.as_mut_ptr(),
224                0,
225                self.n() * size_of::<Self::Scalar>() * self.poly_count(),
226            );
227        }
228    }
229
230    fn zero_at(&mut self, i: usize, j: usize) {
231        unsafe {
232            std::ptr::write_bytes(
233                self.at_mut_ptr(i, j),
234                0,
235                self.n() * size_of::<Self::Scalar>(),
236            );
237        }
238    }
239}
240
241pub trait ZnxRsh: ZnxLayout + ZnxZero
242where
243    Self: Sized,
244    Self::Scalar: IntegerType,
245{
246    fn rsh(&mut self, k: usize, log_base2k: usize, col: usize, carry: &mut [u8]) {
247        rsh(k, log_base2k, self, col, carry)
248    }
249}
250
251pub fn rsh<V: ZnxRsh + ZnxZero>(k: usize, log_base2k: usize, a: &mut V, a_col: usize, tmp_bytes: &mut [u8])
252where
253    V::Scalar: IntegerType,
254{
255    let n: usize = a.n();
256    let size: usize = a.size();
257    let cols: usize = a.cols();
258
259    #[cfg(debug_assertions)]
260    {
261        assert!(
262            tmp_bytes.len() >= rsh_tmp_bytes::<V::Scalar>(n),
263            "invalid carry: carry.len()/size_ofSelf::Scalar={} < rsh_tmp_bytes({}, {})",
264            tmp_bytes.len() / size_of::<V::Scalar>(),
265            n,
266            size,
267        );
268        assert_alignement(tmp_bytes.as_ptr());
269    }
270
271    let size: usize = a.size();
272    let steps: usize = k / log_base2k;
273
274    a.raw_mut().rotate_right(n * steps * cols);
275    (0..cols).for_each(|i| {
276        (0..steps).for_each(|j| {
277            a.zero_at(i, j);
278        })
279    });
280
281    let k_rem: usize = k % log_base2k;
282
283    if k_rem != 0 {
284        let carry: &mut [V::Scalar] = cast_mut(tmp_bytes);
285
286        unsafe {
287            std::ptr::write_bytes(carry.as_mut_ptr(), 0, n * size_of::<V::Scalar>());
288        }
289
290        let log_base2k_t: V::Scalar = V::Scalar::try_from(log_base2k).unwrap();
291        let shift: V::Scalar = V::Scalar::try_from(V::Scalar::BITS as usize - k_rem).unwrap();
292        let k_rem_t: V::Scalar = V::Scalar::try_from(k_rem).unwrap();
293
294        (steps..size).for_each(|i| {
295            izip!(carry.iter_mut(), a.at_mut(a_col, i).iter_mut()).for_each(|(ci, xi)| {
296                *xi += *ci << log_base2k_t;
297                *ci = get_base_k_carry(*xi, shift);
298                *xi = (*xi - *ci) >> k_rem_t;
299            });
300        })
301    }
302}
303
304#[inline(always)]
305fn get_base_k_carry<T: IntegerType>(x: T, shift: T) -> T {
306    (x << shift) >> shift
307}
308
309pub fn rsh_tmp_bytes<T: IntegerType>(n: usize) -> usize {
310    n * std::mem::size_of::<T>()
311}
312
313pub fn switch_degree<T: ZnxLayout + ZnxZero>(b: &mut T, col_b: usize, a: &T, col_a: usize)
314where
315    <T as ZnxLayout>::Scalar: IntegerType,
316{
317    let (n_in, n_out) = (a.n(), b.n());
318    let (gap_in, gap_out): (usize, usize);
319
320    if n_in > n_out {
321        (gap_in, gap_out) = (n_in / n_out, 1)
322    } else {
323        (gap_in, gap_out) = (1, n_out / n_in);
324        b.zero();
325    }
326
327    let size: usize = min(a.size(), b.size());
328
329    (0..size).for_each(|i| {
330        izip!(
331            a.at(col_a, i).iter().step_by(gap_in),
332            b.at_mut(col_b, i).iter_mut().step_by(gap_out)
333        )
334        .for_each(|(x_in, x_out)| *x_out = *x_in);
335    });
336}