1use base2k::{
  2    Backend, MatZnxDft, MatZnxDftAlloc, MatZnxDftToMut, MatZnxDftToRef, Module, VecZnx, VecZnxAlloc, VecZnxDft, VecZnxDftAlloc,
  3    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_k()));
 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 bit precision of the ciphertext.
 47    fn log_k(&self) -> usize;
 48}
 49
 50pub struct RLWECt<C> {
 51    pub data: VecZnx<C>,
 52    pub log_base2k: usize,
 53    pub log_k: usize,
 54}
 55
 56impl<T> Infos for RLWECt<T> {
 57    type Inner = VecZnx<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_k(&self) -> usize {
 68        self.log_k
 69    }
 70}
 71
 72impl<C> VecZnxToMut for RLWECt<C>
 73where
 74    VecZnx<C>: VecZnxToMut,
 75{
 76    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
 77        self.data.to_mut()
 78    }
 79}
 80
 81impl<C> VecZnxToRef for RLWECt<C>
 82where
 83    VecZnx<C>: VecZnxToRef,
 84{
 85    fn to_ref(&self) -> VecZnx<&[u8]> {
 86        self.data.to_ref()
 87    }
 88}
 89
 90pub struct RLWEPt<C> {
 91    pub data: VecZnx<C>,
 92    pub log_base2k: usize,
 93    pub log_k: usize,
 94}
 95
 96impl<T> Infos for RLWEPt<T> {
 97    type Inner = VecZnx<T>;
 98
 99    fn inner(&self) -> &Self::Inner {
100        &self.data
101    }
102
103    fn log_base2k(&self) -> usize {
104        self.log_base2k
105    }
106
107    fn log_k(&self) -> usize {
108        self.log_k
109    }
110}
111
112impl<C> VecZnxToMut for RLWEPt<C>
113where
114    VecZnx<C>: VecZnxToMut,
115{
116    fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
117        self.data.to_mut()
118    }
119}
120
121impl<C> VecZnxToRef for RLWEPt<C>
122where
123    VecZnx<C>: VecZnxToRef,
124{
125    fn to_ref(&self) -> VecZnx<&[u8]> {
126        self.data.to_ref()
127    }
128}
129
130impl RLWECt<Vec<u8>> {
131    pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_k: usize, cols: usize) -> Self {
132        Self {
133            data: module.new_vec_znx(cols, derive_size(log_base2k, log_k)),
134            log_base2k: log_base2k,
135            log_k: log_k,
136        }
137    }
138}
139
140impl RLWEPt<Vec<u8>> {
141    pub fn new<B: Backend>(module: &Module<B>, log_base2k: usize, log_k: usize) -> Self {
142        Self {
143            data: module.new_vec_znx(1, derive_size(log_base2k, log_k)),
144            log_base2k: log_base2k,
145            log_k: log_k,
146        }
147    }
148}
149
150pub struct RLWECtDft<C, B: Backend> {
151    pub data: VecZnxDft<C, B>,
152    pub log_base2k: usize,
153    pub log_k: usize,
154}
155
156impl<B: Backend> RLWECtDft<Vec<u8>, B> {
157    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, cols: usize) -> Self {
158        Self {
159            data: module.new_vec_znx_dft(cols, derive_size(log_base2k, log_k)),
160            log_base2k: log_base2k,
161            log_k: log_k,
162        }
163    }
164}
165
166impl<T, B: Backend> Infos for RLWECtDft<T, B> {
167    type Inner = VecZnxDft<T, B>;
168
169    fn inner(&self) -> &Self::Inner {
170        &self.data
171    }
172
173    fn log_base2k(&self) -> usize {
174        self.log_base2k
175    }
176
177    fn log_k(&self) -> usize {
178        self.log_k
179    }
180}
181
182impl<C, B: Backend> VecZnxDftToMut<B> for RLWECtDft<C, B>
183where
184    VecZnxDft<C, B>: VecZnxDftToMut<B>,
185{
186    fn to_mut(&mut self) -> VecZnxDft<&mut [u8], B> {
187        self.data.to_mut()
188    }
189}
190
191impl<C, B: Backend> VecZnxDftToRef<B> for RLWECtDft<C, B>
192where
193    VecZnxDft<C, B>: VecZnxDftToRef<B>,
194{
195    fn to_ref(&self) -> VecZnxDft<&[u8], B> {
196        self.data.to_ref()
197    }
198}
199
200pub struct GRLWECt<C, B: Backend> {
201    pub data: MatZnxDft<C, B>,
202    pub log_base2k: usize,
203    pub log_k: usize,
204}
205
206impl<B: Backend> GRLWECt<Vec<u8>, B> {
207    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rows: usize) -> Self {
208        Self {
209            data: module.new_mat_znx_dft(rows, 1, 2, derive_size(log_base2k, log_k)),
210            log_base2k: log_base2k,
211            log_k: log_k,
212        }
213    }
214}
215
216impl<T, B: Backend> Infos for GRLWECt<T, B> {
217    type Inner = MatZnxDft<T, B>;
218
219    fn inner(&self) -> &Self::Inner {
220        &self.data
221    }
222
223    fn log_base2k(&self) -> usize {
224        self.log_base2k
225    }
226
227    fn log_k(&self) -> usize {
228        self.log_k
229    }
230}
231
232impl<C, B: Backend> MatZnxDftToMut<B> for GRLWECt<C, B>
233where
234    MatZnxDft<C, B>: MatZnxDftToMut<B>,
235{
236    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
237        self.data.to_mut()
238    }
239}
240
241impl<C, B: Backend> MatZnxDftToRef<B> for GRLWECt<C, B>
242where
243    MatZnxDft<C, B>: MatZnxDftToRef<B>,
244{
245    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
246        self.data.to_ref()
247    }
248}
249
250pub struct RGSWCt<C, B: Backend> {
251    pub data: MatZnxDft<C, B>,
252    pub log_base2k: usize,
253    pub log_k: usize,
254}
255
256impl<B: Backend> RGSWCt<Vec<u8>, B> {
257    pub fn new(module: &Module<B>, log_base2k: usize, log_k: usize, rows: usize) -> Self {
258        Self {
259            data: module.new_mat_znx_dft(rows, 2, 2, derive_size(log_base2k, log_k)),
260            log_base2k: log_base2k,
261            log_k: log_k,
262        }
263    }
264}
265
266impl<T, B: Backend> Infos for RGSWCt<T, B> {
267    type Inner = MatZnxDft<T, B>;
268
269    fn inner(&self) -> &Self::Inner {
270        &self.data
271    }
272
273    fn log_base2k(&self) -> usize {
274        self.log_base2k
275    }
276
277    fn log_k(&self) -> usize {
278        self.log_k
279    }
280}
281
282impl<C, B: Backend> MatZnxDftToMut<B> for RGSWCt<C, B>
283where
284    MatZnxDft<C, B>: MatZnxDftToMut<B>,
285{
286    fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
287        self.data.to_mut()
288    }
289}
290
291impl<C, B: Backend> MatZnxDftToRef<B> for RGSWCt<C, B>
292where
293    MatZnxDft<C, B>: MatZnxDftToRef<B>,
294{
295    fn to_ref(&self) -> MatZnxDft<&[u8], B> {
296        self.data.to_ref()
297    }
298}
299
300pub(crate) fn derive_size(log_base2k: usize, log_k: usize) -> usize {
301    (log_k + log_base2k - 1) / log_base2k
302}