1use base2k::{
  2    Backend, Module, VecZnx, VecZnxAlloc, VecZnxDft, VecZnxDftAlloc, VecZnxDftToMut, VecZnxDftToRef, VecZnxToMut, VecZnxToRef,
  3    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(crate) fn derive_size(log_base2k: usize, log_k: usize) -> usize {
201    (log_k + log_base2k - 1) / log_base2k
202}