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