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