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