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