1use crate::hal::{
  2    layouts::{Data, DataMut, DataRef},
  3    source::Source,
  4};
  5use rand_distr::num_traits::Zero;
  6
  7pub trait ZnxInfos {
  8    /// Returns the ring degree of the polynomials.
  9    fn n(&self) -> usize;
 10
 11    /// Returns the base two logarithm of the ring dimension of the polynomials.
 12    fn log_n(&self) -> usize {
 13        (usize::BITS - (self.n() - 1).leading_zeros()) as _
 14    }
 15
 16    /// Returns the number of rows.
 17    fn rows(&self) -> usize;
 18
 19    /// Returns the number of polynomials in each row.
 20    fn cols(&self) -> usize;
 21
 22    /// Returns the number of size per polynomial.
 23    fn size(&self) -> usize;
 24
 25    /// Returns the total number of small polynomials.
 26    fn poly_count(&self) -> usize {
 27        self.rows() * self.cols() * self.size()
 28    }
 29}
 30
 31pub trait ZnxSliceSize {
 32    /// Returns the slice size, which is the offset between
 33    /// two size of the same column.
 34    fn sl(&self) -> usize;
 35}
 36
 37pub trait DataView {
 38    type D: Data;
 39    fn data(&self) -> &Self::D;
 40}
 41
 42pub trait DataViewMut: DataView {
 43    fn data_mut(&mut self) -> &mut Self::D;
 44}
 45
 46pub trait ZnxView: ZnxInfos + DataView<D: DataRef> {
 47    type Scalar: Copy + Zero;
 48
 49    /// Returns a non-mutable pointer to the underlying coefficients array.
 50    fn as_ptr(&self) -> *const Self::Scalar {
 51        self.data().as_ref().as_ptr() as *const Self::Scalar
 52    }
 53
 54    /// Returns a non-mutable reference to the entire underlying coefficient array.
 55    fn raw(&self) -> &[Self::Scalar] {
 56        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.n() * self.poly_count()) }
 57    }
 58
 59    /// Returns a non-mutable pointer starting at the j-th small polynomial of the i-th column.
 60    fn at_ptr(&self, i: usize, j: usize) -> *const Self::Scalar {
 61        #[cfg(debug_assertions)]
 62        {
 63            assert!(i < self.cols(), "cols: {} >= {}", i, self.cols());
 64            assert!(j < self.size(), "size: {} >= {}", j, self.size());
 65        }
 66        let offset: usize = self.n() * (j * self.cols() + i);
 67        unsafe { self.as_ptr().add(offset) }
 68    }
 69
 70    /// Returns non-mutable reference to the (i, j)-th small polynomial.
 71    fn at(&self, i: usize, j: usize) -> &[Self::Scalar] {
 72        unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n()) }
 73    }
 74}
 75
 76pub trait ZnxViewMut: ZnxView + DataViewMut<D: DataMut> {
 77    /// Returns a mutable pointer to the underlying coefficients array.
 78    fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
 79        self.data_mut().as_mut().as_mut_ptr() as *mut Self::Scalar
 80    }
 81
 82    /// Returns a mutable reference to the entire underlying coefficient array.
 83    fn raw_mut(&mut self) -> &mut [Self::Scalar] {
 84        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.n() * self.poly_count()) }
 85    }
 86
 87    /// Returns a mutable pointer starting at the j-th small polynomial of the i-th column.
 88    fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut Self::Scalar {
 89        #[cfg(debug_assertions)]
 90        {
 91            assert!(i < self.cols(), "cols: {} >= {}", i, self.cols());
 92            assert!(j < self.size(), "size: {} >= {}", j, self.size());
 93        }
 94        let offset: usize = self.n() * (j * self.cols() + i);
 95        unsafe { self.as_mut_ptr().add(offset) }
 96    }
 97
 98    /// Returns mutable reference to the (i, j)-th small polynomial.
 99    fn at_mut(&mut self, i: usize, j: usize) -> &mut [Self::Scalar] {
100        unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n()) }
101    }
102}
103
104//(Jay)Note: Can't provide blanket impl. of ZnxView because Scalar is not known
105impl<T> ZnxViewMut for T where T: ZnxView + DataViewMut<D: DataMut> {}
106
107pub trait ZnxZero
108where
109    Self: Sized,
110{
111    fn zero(&mut self);
112    fn zero_at(&mut self, i: usize, j: usize);
113}
114
115pub trait FillUniform {
116    fn fill_uniform(&mut self, source: &mut Source);
117}
118
119pub trait Reset {
120    fn reset(&mut self);
121}