arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/backend/src/hal/api/znx_base.rs
backend/src/hal/api/znx_base.rsRaw
  1use crate::hal::layouts::{Data, DataMut, DataRef};
  2use rand_distr::num_traits::Zero;
  3use sampling::source::Source;
  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
 29pub trait ZnxSliceSize {
 30    /// Returns the slice size, which is the offset between
 31    /// two size of the same column.
 32    fn sl(&self) -> usize;
 33}
 34
 35pub trait DataView {
 36    type D: Data;
 37    fn data(&self) -> &Self::D;
 38}
 39
 40pub trait DataViewMut: DataView {
 41    fn data_mut(&mut self) -> &mut Self::D;
 42}
 43
 44pub trait ZnxView: ZnxInfos + DataView<D: DataRef> {
 45    type Scalar: Copy + Zero;
 46
 47    /// Returns a non-mutable pointer to the underlying coefficients array.
 48    fn as_ptr(&self) -> *const Self::Scalar {
 49        self.data().as_ref().as_ptr() as *const Self::Scalar
 50    }
 51
 52    /// Returns a non-mutable reference to the entire underlying coefficient array.
 53    fn raw(&self) -> &[Self::Scalar] {
 54        unsafe { std::slice::from_raw_parts(self.as_ptr(), self.n() * self.poly_count()) }
 55    }
 56
 57    /// Returns a non-mutable pointer starting at the j-th small polynomial of the i-th column.
 58    fn at_ptr(&self, i: usize, j: usize) -> *const Self::Scalar {
 59        #[cfg(debug_assertions)]
 60        {
 61            assert!(i < self.cols(), "cols: {} >= {}", i, self.cols());
 62            assert!(j < self.size(), "size: {} >= {}", j, self.size());
 63        }
 64        let offset: usize = self.n() * (j * self.cols() + i);
 65        unsafe { self.as_ptr().add(offset) }
 66    }
 67
 68    /// Returns non-mutable reference to the (i, j)-th small polynomial.
 69    fn at(&self, i: usize, j: usize) -> &[Self::Scalar] {
 70        unsafe { std::slice::from_raw_parts(self.at_ptr(i, j), self.n()) }
 71    }
 72}
 73
 74pub trait ZnxViewMut: ZnxView + DataViewMut<D: DataMut> {
 75    /// Returns a mutable pointer to the underlying coefficients array.
 76    fn as_mut_ptr(&mut self) -> *mut Self::Scalar {
 77        self.data_mut().as_mut().as_mut_ptr() as *mut Self::Scalar
 78    }
 79
 80    /// Returns a mutable reference to the entire underlying coefficient array.
 81    fn raw_mut(&mut self) -> &mut [Self::Scalar] {
 82        unsafe { std::slice::from_raw_parts_mut(self.as_mut_ptr(), self.n() * self.poly_count()) }
 83    }
 84
 85    /// Returns a mutable pointer starting at the j-th small polynomial of the i-th column.
 86    fn at_mut_ptr(&mut self, i: usize, j: usize) -> *mut Self::Scalar {
 87        #[cfg(debug_assertions)]
 88        {
 89            assert!(i < self.cols(), "cols: {} >= {}", i, self.cols());
 90            assert!(j < self.size(), "size: {} >= {}", j, self.size());
 91        }
 92        let offset: usize = self.n() * (j * self.cols() + i);
 93        unsafe { self.as_mut_ptr().add(offset) }
 94    }
 95
 96    /// Returns mutable reference to the (i, j)-th small polynomial.
 97    fn at_mut(&mut self, i: usize, j: usize) -> &mut [Self::Scalar] {
 98        unsafe { std::slice::from_raw_parts_mut(self.at_mut_ptr(i, j), self.n()) }
 99    }
100}
101
102//(Jay)Note: Can't provide blanket impl. of ZnxView because Scalar is not known
103impl<T> ZnxViewMut for T where T: ZnxView + DataViewMut<D: DataMut> {}
104
105pub trait ZnxZero
106where
107    Self: Sized,
108{
109    fn zero(&mut self);
110    fn zero_at(&mut self, i: usize, j: usize);
111}
112
113pub trait FillUniform {
114    fn fill_uniform(&mut self, source: &mut Source);
115}
116
117pub trait Reset {
118    fn reset(&mut self);
119}