1use base2k::{
  2    Backend, DataView, DataViewMut, MatZnxDft, MatZnxDftAlloc, MatZnxDftToMut, MatZnxDftToRef, Module, VecZnx, VecZnxAlloc,
  3    VecZnxDft, VecZnxDftAlloc, VecZnxDftToMut, VecZnxDftToRef, VecZnxToMut, VecZnxToRef, ZnxInfos,
  4};
  5
  6pub trait Infos<T>
  7where
  8    T: ZnxInfos,
  9{
 10    fn inner(&self) -> &T;
 11
 12    /// Returns the ring degree of the polynomials.
 13    fn n(&self) -> usize {
 14        self.inner().n()
 15    }
 16
 17    /// Returns the base two logarithm of the ring dimension of the polynomials.
 18    fn log_n(&self) -> usize {
 19        self.inner().log_n()
 20    }
 21
 22    /// Returns the number of rows.
 23    fn rows(&self) -> usize {
 24        self.inner().rows()
 25    }
 26
 27    /// Returns the number of polynomials in each row.
 28    fn cols(&self) -> usize {
 29        self.inner().cols()
 30    }
 31
 32    /// Returns the number of size per polynomial.
 33    fn size(&self) -> usize {
 34        let size: usize = self.inner().size();
 35        debug_assert_eq!(size, derive_size(self.log_base2k(), self.log_q()));
 36        size
 37    }
 38
 39    /// Returns the total number of small polynomials.
 40    fn poly_count(&self) -> usize {
 41        self.rows() * self.cols() * self.size()
 42    }
 43
 44    /// Returns the base 2 logarithm of the ciphertext base.
 45    fn log_base2k(&self) -> usize;
 46
 47    /// Returns the base 2 logarithm of the ciphertext modulus.
 48    fn log_q(&self) -> usize;
 49}
 50
 51pub struct Ciphertext<T> {
 52    data: T,
 53    log_base2k: usize,
 54    log_q: usize,
 55}
 56
 57impl<T> Infos<T> for Ciphertext<T>
 58where
 59    T: ZnxInfos,
 60{
 61    fn inner(&self) -> &T {
 62        &self.data
 63    }
 64
 65    fn log_base2k(&self) -> usize {
 66        self.log_base2k
 67    }
 68
 69    fn log_q(&self) -> usize {
 70        self.log_q
 71    }
 72}
 73
 74impl<D> DataView for Ciphertext<D> {
 75    type D = D;
 76    fn data(&self) -> &Self::D {
 77        &self.data
 78    }
 79}
 80
 81impl<D> DataViewMut for Ciphertext<D> {
 82    fn data_mut(&mut self) -> &mut Self::D {
 83        &mut self.data
 84    }
 85}
 86
 87pub struct Plaintext<T> {
 88    data: T,
 89    log_base2k: usize,
 90    log_q: usize,
 91}
 92
 93impl<T> Infos<T> for Plaintext<T>
 94where
 95    T: ZnxInfos,
 96{
 97    fn inner(&self) -> &T {
 98        &self.data
 99    }
100
101    fn log_base2k(&self) -> usize {
102        self.log_base2k
103    }
104
105    fn log_q(&self) -> usize {
106        self.log_q
107    }
108}
109
110impl<T> Plaintext<T> {
111    pub fn data(&self) -> &T {
112        &self.data
113    }
114
115    pub fn data_mut(&mut self) -> &mut T {
116        &mut self.data
117    }
118}
119
120pub(crate) type CtVecZnx<C> = Ciphertext<VecZnx<C>>;
121pub(crate) type CtVecZnxDft<C, B: Backend> = Ciphertext<VecZnxDft<C, B>>;
122pub(crate) type CtMatZnxDft<C, B: Backend> = Ciphertext<MatZnxDft<C, B>>;
123pub(crate) type PtVecZnx<C> = Plaintext<VecZnx<C>>;
124pub(crate) type PtVecZnxDft<C, B: Backend> = Plaintext<VecZnxDft<C, B>>;
125pub(crate) type PtMatZnxDft<C, B: Backend> = Plaintext<MatZnxDft<C, B>>;
126
127impl<D> VecZnxToMut for Ciphertext<D>
128where
129    D: VecZnxToMut,
130{
131    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
132        self.data_mut().to_mut()
133    }
134}
135
136impl<D> VecZnxToRef for Ciphertext<D>
137where
138    D: VecZnxToRef,
139{
140    fn to_ref(&self) -> VecZnx<&[u8]> {
141        self.data().to_ref()
142    }
143}
144
145impl Ciphertext<VecZnx<Vec<u8>>> {
146    pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_q: usize, cols: usize) -> Self {
147        Self {
148            data: module.new_vec_znx(cols, derive_size(log_base2k, log_q)),
149            log_base2k: log_base2k,
150            log_q: log_q,
151        }
152    }
153}
154
155impl<D> VecZnxToMut for Plaintext<D>
156where
157    D: VecZnxToMut,
158{
159    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
160        self.data_mut().to_mut()
161    }
162}
163
164impl<D> VecZnxToRef for Plaintext<D>
165where
166    D: VecZnxToRef,
167{
168    fn to_ref(&self) -> VecZnx<&[u8]> {
169        self.data().to_ref()
170    }
171}
172
173impl Plaintext<VecZnx<Vec<u8>>> {
174    pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_q: usize) -> Self {
175        Self {
176            data: module.new_vec_znx(1, derive_size(log_base2k, log_q)),
177            log_base2k: log_base2k,
178            log_q: log_q,
179        }
180    }
181}
182
183impl<D, B: Backend> VecZnxDftToMut<B> for Ciphertext<D>
184where
185    D: VecZnxDftToMut<B>,
186{
187    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
188        self.data_mut().to_mut()
189    }
190}
191
192impl<D, B: Backend> VecZnxDftToRef<B> for Ciphertext<D>
193where
194    D: VecZnxDftToRef<B>,
195{
196    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
197        self.data().to_ref()
198    }
199}
200
201impl<B: Backend> Ciphertext<VecZnxDft<Vec<u8>, B>> {
202    pub fn new(module: &Module<B>, log_base2k: usize, log_q: usize, cols: usize) -> Self {
203        Self {
204            data: module.new_vec_znx_dft(cols, derive_size(log_base2k, log_q)),
205            log_base2k: log_base2k,
206            log_q: log_q,
207        }
208    }
209}
210
211impl<D, B: Backend> MatZnxDftToMut<B> for Ciphertext<D>
212where
213    D: MatZnxDftToMut<B>,
214{
215    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
216        self.data_mut().to_mut()
217    }
218}
219
220impl<D, B: Backend> MatZnxDftToRef<B> for Ciphertext<D>
221where
222    D: MatZnxDftToRef<B>,
223{
224    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
225        self.data().to_ref()
226    }
227}
228
229impl<B: Backend> Ciphertext<MatZnxDft<Vec<u8>, B>> {
230    pub fn new(module: &Module<B>, log_base2k: usize, rows: usize, cols_in: usize, cols_out: usize, log_q: usize) -> Self {
231        Self {
232            data: module.new_mat_znx_dft(rows, cols_in, cols_out, derive_size(log_base2k, log_q)),
233            log_base2k: log_base2k,
234            log_q: log_q,
235        }
236    }
237}
238
239pub(crate) fn derive_size(log_base2k: usize, log_q: usize) -> usize {
240    (log_q + log_base2k - 1) / log_base2k
241}