1use backend::hal::{
  2    api::{MatZnxAlloc, MatZnxAllocBytes, VmpPMatAlloc, VmpPMatAllocBytes, VmpPMatPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, Scratch, VmpPMat, WriterTo},
  4};
  5
  6use crate::{GLWECiphertext, Infos};
  7
  8pub trait GGLWEExecLayoutFamily<B: Backend> = VmpPMatAlloc<B> + VmpPMatAllocBytes + VmpPMatPrepare<B>;
  9
 10#[derive(PartialEq, Eq)]
 11pub struct GGLWECiphertext<D: Data> {
 12    pub(crate) data: MatZnx<D>,
 13    pub(crate) basek: usize,
 14    pub(crate) k: usize,
 15    pub(crate) digits: usize,
 16}
 17
 18impl<D: DataRef> GGLWECiphertext<D> {
 19    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
 20        GLWECiphertext {
 21            data: self.data.at(row, col),
 22            basek: self.basek,
 23            k: self.k,
 24        }
 25    }
 26}
 27
 28impl<D: DataMut> GGLWECiphertext<D> {
 29    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
 30        GLWECiphertext {
 31            data: self.data.at_mut(row, col),
 32            basek: self.basek,
 33            k: self.k,
 34        }
 35    }
 36}
 37
 38impl GGLWECiphertext<Vec<u8>> {
 39    pub fn alloc<B: Backend>(
 40        module: &Module<B>,
 41        basek: usize,
 42        k: usize,
 43        rows: usize,
 44        digits: usize,
 45        rank_in: usize,
 46        rank_out: usize,
 47    ) -> Self
 48    where
 49        Module<B>: MatZnxAlloc,
 50    {
 51        let size: usize = k.div_ceil(basek);
 52        debug_assert!(
 53            size > digits,
 54            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
 55            size,
 56            digits
 57        );
 58
 59        assert!(
 60            rows * digits <= size,
 61            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
 62            rows,
 63            digits,
 64            size
 65        );
 66
 67        Self {
 68            data: module.mat_znx_alloc(rows, rank_in, rank_out + 1, size),
 69            basek: basek,
 70            k,
 71            digits,
 72        }
 73    }
 74
 75    pub fn bytes_of<B: Backend>(
 76        module: &Module<B>,
 77        basek: usize,
 78        k: usize,
 79        rows: usize,
 80        digits: usize,
 81        rank_in: usize,
 82        rank_out: usize,
 83    ) -> usize
 84    where
 85        Module<B>: MatZnxAllocBytes,
 86    {
 87        let size: usize = k.div_ceil(basek);
 88        debug_assert!(
 89            size > digits,
 90            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
 91            size,
 92            digits
 93        );
 94
 95        assert!(
 96            rows * digits <= size,
 97            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
 98            rows,
 99            digits,
100            size
101        );
102
103        module.mat_znx_alloc_bytes(rows, rank_in, rank_out + 1, rows)
104    }
105}
106
107impl<D: Data> Infos for GGLWECiphertext<D> {
108    type Inner = MatZnx<D>;
109
110    fn inner(&self) -> &Self::Inner {
111        &self.data
112    }
113
114    fn basek(&self) -> usize {
115        self.basek
116    }
117
118    fn k(&self) -> usize {
119        self.k
120    }
121}
122
123impl<D: Data> GGLWECiphertext<D> {
124    pub fn rank(&self) -> usize {
125        self.data.cols_out() - 1
126    }
127
128    pub fn digits(&self) -> usize {
129        self.digits
130    }
131
132    pub fn rank_in(&self) -> usize {
133        self.data.cols_in()
134    }
135
136    pub fn rank_out(&self) -> usize {
137        self.data.cols_out() - 1
138    }
139}
140
141use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
142
143impl<D: DataMut> ReaderFrom for GGLWECiphertext<D> {
144    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
145        self.k = reader.read_u64::<LittleEndian>()? as usize;
146        self.basek = reader.read_u64::<LittleEndian>()? as usize;
147        self.digits = reader.read_u64::<LittleEndian>()? as usize;
148        self.data.read_from(reader)
149    }
150}
151
152impl<D: DataRef> WriterTo for GGLWECiphertext<D> {
153    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
154        writer.write_u64::<LittleEndian>(self.k as u64)?;
155        writer.write_u64::<LittleEndian>(self.basek as u64)?;
156        writer.write_u64::<LittleEndian>(self.digits as u64)?;
157        self.data.write_to(writer)
158    }
159}
160
161#[derive(PartialEq, Eq)]
162pub struct GGLWECiphertextExec<D: Data, B: Backend> {
163    pub(crate) data: VmpPMat<D, B>,
164    pub(crate) basek: usize,
165    pub(crate) k: usize,
166    pub(crate) digits: usize,
167}
168
169impl<B: Backend> GGLWECiphertextExec<Vec<u8>, B> {
170    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize, rank_out: usize) -> Self
171    where
172        Module<B>: GGLWEExecLayoutFamily<B>,
173    {
174        let size: usize = k.div_ceil(basek);
175        debug_assert!(
176            size > digits,
177            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
178            size,
179            digits
180        );
181
182        assert!(
183            rows * digits <= size,
184            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
185            rows,
186            digits,
187            size
188        );
189
190        Self {
191            data: module.vmp_pmat_alloc(rows, rank_in, rank_out + 1, size),
192            basek: basek,
193            k,
194            digits,
195        }
196    }
197
198    pub fn bytes_of(
199        module: &Module<B>,
200        basek: usize,
201        k: usize,
202        rows: usize,
203        digits: usize,
204        rank_in: usize,
205        rank_out: usize,
206    ) -> usize
207    where
208        Module<B>: GGLWEExecLayoutFamily<B>,
209    {
210        let size: usize = k.div_ceil(basek);
211        debug_assert!(
212            size > digits,
213            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
214            size,
215            digits
216        );
217
218        assert!(
219            rows * digits <= size,
220            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
221            rows,
222            digits,
223            size
224        );
225
226        module.vmp_pmat_alloc_bytes(rows, rank_in, rank_out + 1, rows)
227    }
228}
229
230impl<D: Data, B: Backend> Infos for GGLWECiphertextExec<D, B> {
231    type Inner = VmpPMat<D, B>;
232
233    fn inner(&self) -> &Self::Inner {
234        &self.data
235    }
236
237    fn basek(&self) -> usize {
238        self.basek
239    }
240
241    fn k(&self) -> usize {
242        self.k
243    }
244}
245
246impl<D: Data, B: Backend> GGLWECiphertextExec<D, B> {
247    pub fn rank(&self) -> usize {
248        self.data.cols_out() - 1
249    }
250
251    pub fn digits(&self) -> usize {
252        self.digits
253    }
254
255    pub fn rank_in(&self) -> usize {
256        self.data.cols_in()
257    }
258
259    pub fn rank_out(&self) -> usize {
260        self.data.cols_out() - 1
261    }
262}
263
264impl<D: DataMut, B: Backend> GGLWECiphertextExec<D, B> {
265    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GGLWECiphertext<DataOther>, scratch: &mut Scratch<B>)
266    where
267        DataOther: DataRef,
268        Module<B>: GGLWEExecLayoutFamily<B>,
269    {
270        module.vmp_prepare(&mut self.data, &other.data, scratch);
271        self.basek = other.basek;
272        self.k = other.k;
273        self.digits = other.digits;
274    }
275}