1use itertools::izip;
  2use rand_distr::num_traits::Zero;
  3use std::cmp::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}
 32
 33pub trait DataView {
 34    type D;
 35    fn data(&self) -> &Self::D;
 36}
 37
 38pub trait DataViewMut: DataView {
 39    fn data_mut(&mut self) -> &mut Self::D;
 40}
 41
 42pub trait ZnxView: ZnxInfos + DataView<D: AsRef<[u8]>> {
 43    type Scalar: Copy;
 44
 45    /// Returns a non-mutable pointer to the underlying coefficients array.
 46    fn as_ptr(&self) -> *const Self::Scalar {
 47        self.data().as_ref().as_ptr() as *const Self::Scalar
 48    }
 49
 50    /// Returns a non-mutable reference to the entire underlying coefficient array.
 51    fn raw(&self) -> &[Self::Scalar] {
 52        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.n() * self.poly_count()) }
 53    }
 54
 55    /// Returns a non-mutable pointer starting at the j-th small polynomial of the i-th column.
 56    fn at_ptr(&self, i: usize, j: usize) -> *const Self::Scalar {
 57        #[cfg(debug_assertions)]
 58        {
 59            assert!(i < self.cols());
 60            assert!(j < self.size());
 61        }
 62        let offset: usize = self.n() * (j * self.cols() + i);
 63        unsafe { self.as_ptr().add(offset) }
 64    }
 65
 66    /// Returns non-mutable reference to the (i, j)-th small polynomial.
 67    fn at(&self, i: usize, j: usize) -> &[Self::Scalar] {
 68        unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n()) }
 69    }
 70}
 71
 72pub trait ZnxViewMut: ZnxView + DataViewMut<D: AsMut<[u8]>> {
 73    /// Returns a mutable pointer to the underlying coefficients array.
 74    fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
 75        self.data_mut().as_mut().as_mut_ptr() as *mut Self::Scalar
 76    }
 77
 78    /// Returns a mutable reference to the entire underlying coefficient array.
 79    fn raw_mut(&mut self) -> &mut [Self::Scalar] {
 80        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.n() * self.poly_count()) }
 81    }
 82
 83    /// Returns a mutable pointer starting at the j-th small polynomial of the i-th column.
 84    fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut Self::Scalar {
 85        #[cfg(debug_assertions)]
 86        {
 87            assert!(i < self.cols());
 88            assert!(j < self.size());
 89        }
 90        let offset: usize = self.n() * (j * self.cols() + i);
 91        unsafe { self.as_mut_ptr().add(offset) }
 92    }
 93
 94    /// Returns mutable reference to the (i, j)-th small polynomial.
 95    fn at_mut(&mut self, i: usize, j: usize) -> &mut [Self::Scalar] {
 96        unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n()) }
 97    }
 98}
 99
100//(Jay)Note: Can't provide blanket impl. of ZnxView because Scalar is not known
101impl<T> ZnxViewMut for T where T: ZnxView + DataViewMut<D: AsMut<[u8]>> {}
102
103pub trait ZnxZero: ZnxViewMut
104where
105    Self: Sized,
106{
107    fn zero(&mut self) {
108        unsafe {
109            std::ptr::write_bytes(
110                self.as_mut_ptr(),
111                0,
112                self.n() * size_of::<Self::Scalar>() * self.poly_count(),
113            );
114        }
115    }
116
117    fn zero_at(&mut self, i: usize, j: usize) {
118        unsafe {
119            std::ptr::write_bytes(
120                self.at_mut_ptr(i, j),
121                0,
122                self.n() * size_of::<Self::Scalar>(),
123            );
124        }
125    }
126}
127
128// Blanket implementations
129impl<T> ZnxZero for T where T: ZnxViewMut {}
130// impl<T> ZnxRsh for T where T: ZnxZero {}
131
132pub fn switch_degree<S: Copy, DMut: ZnxViewMut<Scalar = S> + ZnxZero, D: ZnxView<Scalar = S>>(
133    b: &mut DMut,
134    col_b: usize,
135    a: &D,
136    col_a: usize,
137) {
138    let (n_in, n_out) = (a.n(), b.n());
139    let (gap_in, gap_out): (usize, usize);
140
141    if n_in > n_out {
142        (gap_in, gap_out) = (n_in / n_out, 1)
143    } else {
144        (gap_in, gap_out) = (1, n_out / n_in);
145        b.zero();
146    }
147
148    let size: usize = min(a.size(), b.size());
149
150    (0..size).for_each(|i| {
151        izip!(
152            a.at(col_a, i).iter().step_by(gap_in),
153            b.at_mut(col_b, i).iter_mut().step_by(gap_out)
154        )
155        .for_each(|(x_in, x_out)| *x_out = *x_in);
156    });
157}
158
159use std::ops::{Add, AddAssign, Div, Mul, Neg, Shl, Shr, Sub};
160
161use crate::Scratch;
162pub trait Integer:
163    Copy
164    + Default
165    + PartialEq
166    + PartialOrd
167    + Add<Output = Self>
168    + Sub<Output = Self>
169    + Mul<Output = Self>
170    + Div<Output = Self>
171    + Neg<Output = Self>
172    + Shl<Output = Self>
173    + Shr<Output = Self>
174    + AddAssign
175{
176    const BITS: u32;
177}
178
179impl Integer for i64 {
180    const BITS: u32 = 64;
181}
182
183impl Integer for i128 {
184    const BITS: u32 = 128;
185}
186
187//(Jay)Note: `rsh` impl. ignores the column
188pub fn rsh<V: ZnxZero>(k: usize, log_base2k: usize, a: &mut V, _a_col: usize, scratch: &mut Scratch)
189where
190    V::Scalar: From<usize> + Integer + Zero,
191{
192    let n: usize = a.n();
193    let _size: usize = a.size();
194    let cols: usize = a.cols();
195
196    let size: usize = a.size();
197    let steps: usize = k / log_base2k;
198
199    a.raw_mut().rotate_right(n * steps * cols);
200    (0..cols).for_each(|i| {
201        (0..steps).for_each(|j| {
202            a.zero_at(i, j);
203        })
204    });
205
206    let k_rem: usize = k % log_base2k;
207
208    if k_rem != 0 {
209        let (carry, _) = scratch.tmp_scalar_slice::<V::Scalar>(rsh_tmp_bytes::<V::Scalar>(n));
210
211        unsafe {
212            std::ptr::write_bytes(carry.as_mut_ptr(), 0, n * size_of::<V::Scalar>());
213        }
214
215        let log_base2k_t = V::Scalar::from(log_base2k);
216        let shift = V::Scalar::from(V::Scalar::BITS as usize - k_rem);
217        let k_rem_t = V::Scalar::from(k_rem);
218
219        (0..cols).for_each(|i| {
220            (steps..size).for_each(|j| {
221                izip!(carry.iter_mut(), a.at_mut(i, j).iter_mut()).for_each(|(ci, xi)| {
222                    *xi += *ci << log_base2k_t;
223                    *ci = (*xi << shift) >> shift;
224                    *xi = (*xi - *ci) >> k_rem_t;
225                });
226            });
227            carry.iter_mut().for_each(|r| *r = V::Scalar::zero());
228        })
229    }
230}
231
232pub fn rsh_tmp_bytes<T>(n: usize) -> usize {
233    n * std::mem::size_of::<T>()
234}