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