1use itertools::izip;
  2use rand_distr::num_traits::Zero;
  3
  4pub trait ZnxInfos {
  5    /// Returns the ring degree of the polynomials.
  6    fn n(&self) -> usize;
  7
  8    /// Returns the base two logarithm of the ring dimension of the polynomials.
  9    fn log_n(&self) -> usize {
 10        (usize::BITS - (self.n() - 1).leading_zeros()) as _
 11    }
 12
 13    /// Returns the number of rows.
 14    fn rows(&self) -> usize;
 15
 16    /// Returns the number of polynomials in each row.
 17    fn cols(&self) -> usize;
 18
 19    /// Returns the number of size per polynomial.
 20    fn size(&self) -> usize;
 21
 22    /// Returns the total number of small polynomials.
 23    fn poly_count(&self) -> usize {
 24        self.rows() * self.cols() * self.size()
 25    }
 26}
 27
 28pub trait ZnxSliceSize {
 29    /// Returns the slice size, which is the offset between
 30    /// two size of the same column.
 31    fn sl(&self) -> usize;
 32}
 33
 34pub trait DataView {
 35    type D;
 36    fn data(&self) -> &Self::D;
 37}
 38
 39pub trait DataViewMut: DataView {
 40    fn data_mut(&mut self) -> &mut Self::D;
 41}
 42
 43pub trait ZnxView: ZnxInfos + DataView<D: AsRef<[u8]>> {
 44    type Scalar: Copy;
 45
 46    /// Returns a non-mutable pointer to the underlying coefficients array.
 47    fn as_ptr(&self) -> *const Self::Scalar {
 48        self.data().as_ref().as_ptr() as *const Self::Scalar
 49    }
 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 non-mutable pointer starting at the j-th small polynomial of the i-th column.
 57    fn at_ptr(&self, i: usize, j: usize) -> *const Self::Scalar {
 58        #[cfg(debug_assertions)]
 59        {
 60            assert!(i < self.cols());
 61            assert!(j < self.size());
 62        }
 63        let offset: usize = self.n() * (j * self.cols() + i);
 64        unsafe { self.as_ptr().add(offset) }
 65    }
 66
 67    /// Returns non-mutable reference to the (i, j)-th small polynomial.
 68    fn at(&self, i: usize, j: usize) -> &[Self::Scalar] {
 69        unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n()) }
 70    }
 71}
 72
 73pub trait ZnxViewMut: ZnxView + DataViewMut<D: AsMut<[u8]>> {
 74    /// Returns a mutable pointer to the underlying coefficients array.
 75    fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
 76        self.data_mut().as_mut().as_mut_ptr() as *mut Self::Scalar
 77    }
 78
 79    /// Returns a mutable reference to the entire underlying coefficient array.
 80    fn raw_mut(&mut self) -> &mut [Self::Scalar] {
 81        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.n() * self.poly_count()) }
 82    }
 83
 84    /// Returns a mutable pointer starting at the j-th small polynomial of the i-th column.
 85    fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut Self::Scalar {
 86        #[cfg(debug_assertions)]
 87        {
 88            assert!(i < self.cols());
 89            assert!(j < self.size());
 90        }
 91        let offset: usize = self.n() * (j * self.cols() + i);
 92        unsafe { self.as_mut_ptr().add(offset) }
 93    }
 94
 95    /// Returns mutable reference to the (i, j)-th small polynomial.
 96    fn at_mut(&mut self, i: usize, j: usize) -> &mut [Self::Scalar] {
 97        unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n()) }
 98    }
 99}
100
101//(Jay)Note: Can't provide blanket impl. of ZnxView because Scalar is not known
102impl<T> ZnxViewMut for T where T: ZnxView + DataViewMut<D: AsMut<[u8]>> {}
103
104pub trait ZnxZero: ZnxViewMut
105where
106    Self: Sized,
107{
108    fn zero(&mut self) {
109        unsafe {
110            std::ptr::write_bytes(
111                self.as_mut_ptr(),
112                0,
113                self.n() * size_of::<Self::Scalar>() * self.poly_count(),
114            );
115        }
116    }
117
118    fn zero_at(&mut self, i: usize, j: usize) {
119        unsafe {
120            std::ptr::write_bytes(
121                self.at_mut_ptr(i, j),
122                0,
123                self.n() * size_of::<Self::Scalar>(),
124            );
125        }
126    }
127}
128
129// Blanket implementations
130impl<T> ZnxZero for T where T: ZnxViewMut {}
131// impl<T> ZnxRsh for T where T: ZnxZero {}
132
133use std::ops::{Add, AddAssign, Div, Mul, Neg, Shl, Shr, Sub};
134
135use crate::Scratch;
136pub trait Integer:
137    Copy
138    + Default
139    + PartialEq
140    + PartialOrd
141    + Add<Output = Self>
142    + Sub<Output = Self>
143    + Mul<Output = Self>
144    + Div<Output = Self>
145    + Neg<Output = Self>
146    + Shl<Output = Self>
147    + Shr<Output = Self>
148    + AddAssign
149{
150    const BITS: u32;
151}
152
153impl Integer for i64 {
154    const BITS: u32 = 64;
155}
156
157impl Integer for i128 {
158    const BITS: u32 = 128;
159}
160
161//(Jay)Note: `rsh` impl. ignores the column
162pub fn rsh<V: ZnxZero>(k: usize, log_base2k: usize, a: &mut V, _a_col: usize, scratch: &mut Scratch)
163where
164    V::Scalar: From<usize> + Integer + Zero,
165{
166    let n: usize = a.n();
167    let _size: usize = a.size();
168    let cols: usize = a.cols();
169
170    let size: usize = a.size();
171    let steps: usize = k / log_base2k;
172
173    a.raw_mut().rotate_right(n * steps * cols);
174    (0..cols).for_each(|i| {
175        (0..steps).for_each(|j| {
176            a.zero_at(i, j);
177        })
178    });
179
180    let k_rem: usize = k % log_base2k;
181
182    if k_rem != 0 {
183        let (carry, _) = scratch.tmp_scalar_slice::<V::Scalar>(rsh_tmp_bytes::<V::Scalar>(n));
184
185        unsafe {
186            std::ptr::write_bytes(carry.as_mut_ptr(), 0, n * size_of::<V::Scalar>());
187        }
188
189        let log_base2k_t = V::Scalar::from(log_base2k);
190        let shift = V::Scalar::from(V::Scalar::BITS as usize - k_rem);
191        let k_rem_t = V::Scalar::from(k_rem);
192
193        (0..cols).for_each(|i| {
194            (steps..size).for_each(|j| {
195                izip!(carry.iter_mut(), a.at_mut(i, j).iter_mut()).for_each(|(ci, xi)| {
196                    *xi += *ci << log_base2k_t;
197                    *ci = (*xi << shift) >> shift;
198                    *xi = (*xi - *ci) >> k_rem_t;
199                });
200            });
201            carry.iter_mut().for_each(|r| *r = V::Scalar::zero());
202        })
203    }
204}
205
206pub fn rsh_tmp_bytes<T>(n: usize) -> usize {
207    n * std::mem::size_of::<T>()
208}