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}