1use poulpy_hal::{
  2    layouts::{Data, DataMut, DataRef, FillUniform, MatZnx, ReaderFrom, WriterTo, ZnxInfos},
  3    source::Source,
  4};
  5
  6use crate::layouts::{Base2K, BuildError, Degree, Dnum, Dsize, GLWECiphertext, GLWEInfos, LWEInfos, Rank, TorusPrecision};
  7use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
  8
  9use std::fmt;
 10
 11pub trait GGLWEInfos
 12where
 13    Self: GLWEInfos,
 14{
 15    fn dnum(&self) -> Dnum;
 16    fn dsize(&self) -> Dsize;
 17    fn rank_in(&self) -> Rank;
 18    fn rank_out(&self) -> Rank;
 19    fn layout(&self) -> GGLWECiphertextLayout {
 20        GGLWECiphertextLayout {
 21            n: self.n(),
 22            base2k: self.base2k(),
 23            k: self.k(),
 24            rank_in: self.rank_in(),
 25            rank_out: self.rank_out(),
 26            dsize: self.dsize(),
 27            dnum: self.dnum(),
 28        }
 29    }
 30}
 31
 32#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 33pub struct GGLWECiphertextLayout {
 34    pub n: Degree,
 35    pub base2k: Base2K,
 36    pub k: TorusPrecision,
 37    pub rank_in: Rank,
 38    pub rank_out: Rank,
 39    pub dnum: Dnum,
 40    pub dsize: Dsize,
 41}
 42
 43impl LWEInfos for GGLWECiphertextLayout {
 44    fn base2k(&self) -> Base2K {
 45        self.base2k
 46    }
 47
 48    fn k(&self) -> TorusPrecision {
 49        self.k
 50    }
 51
 52    fn n(&self) -> Degree {
 53        self.n
 54    }
 55}
 56
 57impl GLWEInfos for GGLWECiphertextLayout {
 58    fn rank(&self) -> Rank {
 59        self.rank_out
 60    }
 61}
 62
 63impl GGLWEInfos for GGLWECiphertextLayout {
 64    fn rank_in(&self) -> Rank {
 65        self.rank_in
 66    }
 67
 68    fn dsize(&self) -> Dsize {
 69        self.dsize
 70    }
 71
 72    fn rank_out(&self) -> Rank {
 73        self.rank_out
 74    }
 75
 76    fn dnum(&self) -> Dnum {
 77        self.dnum
 78    }
 79}
 80
 81#[derive(PartialEq, Eq, Clone)]
 82pub struct GGLWECiphertext<D: Data> {
 83    pub(crate) data: MatZnx<D>,
 84    pub(crate) k: TorusPrecision,
 85    pub(crate) base2k: Base2K,
 86    pub(crate) dsize: Dsize,
 87}
 88
 89impl<D: Data> LWEInfos for GGLWECiphertext<D> {
 90    fn base2k(&self) -> Base2K {
 91        self.base2k
 92    }
 93
 94    fn k(&self) -> TorusPrecision {
 95        self.k
 96    }
 97
 98    fn n(&self) -> Degree {
 99        Degree(self.data.n() as u32)
100    }
101
102    fn size(&self) -> usize {
103        self.data.size()
104    }
105}
106
107impl<D: Data> GLWEInfos for GGLWECiphertext<D> {
108    fn rank(&self) -> Rank {
109        self.rank_out()
110    }
111}
112
113impl<D: Data> GGLWEInfos for GGLWECiphertext<D> {
114    fn rank_in(&self) -> Rank {
115        Rank(self.data.cols_in() as u32)
116    }
117
118    fn rank_out(&self) -> Rank {
119        Rank(self.data.cols_out() as u32 - 1)
120    }
121
122    fn dsize(&self) -> Dsize {
123        self.dsize
124    }
125
126    fn dnum(&self) -> Dnum {
127        Dnum(self.data.rows() as u32)
128    }
129}
130
131pub struct GGLWECiphertextBuilder<D: Data> {
132    data: Option<MatZnx<D>>,
133    base2k: Option<Base2K>,
134    k: Option<TorusPrecision>,
135    dsize: Option<Dsize>,
136}
137
138impl<D: Data> GGLWECiphertext<D> {
139    #[inline]
140    pub fn builder() -> GGLWECiphertextBuilder<D> {
141        GGLWECiphertextBuilder {
142            data: None,
143            base2k: None,
144            k: None,
145            dsize: None,
146        }
147    }
148}
149
150impl GGLWECiphertextBuilder<Vec<u8>> {
151    #[inline]
152    pub fn layout<A>(mut self, infos: &A) -> Self
153    where
154        A: GGLWEInfos,
155    {
156        self.data = Some(MatZnx::alloc(
157            infos.n().into(),
158            infos.dnum().into(),
159            infos.rank_in().into(),
160            (infos.rank_out() + 1).into(),
161            infos.size(),
162        ));
163        self.base2k = Some(infos.base2k());
164        self.k = Some(infos.k());
165        self.dsize = Some(infos.dsize());
166        self
167    }
168}
169
170impl<D: Data> GGLWECiphertextBuilder<D> {
171    #[inline]
172    pub fn data(mut self, data: MatZnx<D>) -> Self {
173        self.data = Some(data);
174        self
175    }
176    #[inline]
177    pub fn base2k(mut self, base2k: Base2K) -> Self {
178        self.base2k = Some(base2k);
179        self
180    }
181    #[inline]
182    pub fn k(mut self, k: TorusPrecision) -> Self {
183        self.k = Some(k);
184        self
185    }
186
187    #[inline]
188    pub fn dsize(mut self, dsize: Dsize) -> Self {
189        self.dsize = Some(dsize);
190        self
191    }
192
193    pub fn build(self) -> Result<GGLWECiphertext<D>, BuildError> {
194        let data: MatZnx<D> = self.data.ok_or(BuildError::MissingData)?;
195        let base2k: Base2K = self.base2k.ok_or(BuildError::MissingBase2K)?;
196        let k: TorusPrecision = self.k.ok_or(BuildError::MissingK)?;
197        let dsize: Dsize = self.dsize.ok_or(BuildError::MissingDigits)?;
198
199        if base2k == 0_u32 {
200            return Err(BuildError::ZeroBase2K);
201        }
202
203        if dsize == 0_u32 {
204            return Err(BuildError::ZeroBase2K);
205        }
206
207        if k == 0_u32 {
208            return Err(BuildError::ZeroTorusPrecision);
209        }
210
211        if data.n() == 0 {
212            return Err(BuildError::ZeroDegree);
213        }
214
215        if data.cols() == 0 {
216            return Err(BuildError::ZeroCols);
217        }
218
219        if data.size() == 0 {
220            return Err(BuildError::ZeroLimbs);
221        }
222
223        Ok(GGLWECiphertext {
224            data,
225            base2k,
226            k,
227            dsize,
228        })
229    }
230}
231
232impl<D: DataRef> GGLWECiphertext<D> {
233    pub fn data(&self) -> &MatZnx<D> {
234        &self.data
235    }
236}
237
238impl<D: DataMut> GGLWECiphertext<D> {
239    pub fn data_mut(&mut self) -> &mut MatZnx<D> {
240        &mut self.data
241    }
242}
243
244impl<D: DataRef> fmt::Debug for GGLWECiphertext<D> {
245    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
246        write!(f, "{self}")
247    }
248}
249
250impl<D: DataMut> FillUniform for GGLWECiphertext<D> {
251    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
252        self.data.fill_uniform(log_bound, source);
253    }
254}
255
256impl<D: DataRef> fmt::Display for GGLWECiphertext<D> {
257    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
258        write!(
259            f,
260            "(GGLWECiphertext: k={} base2k={} dsize={}) {}",
261            self.k().0,
262            self.base2k().0,
263            self.dsize().0,
264            self.data
265        )
266    }
267}
268
269impl<D: DataRef> GGLWECiphertext<D> {
270    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
271        GLWECiphertext::builder()
272            .data(self.data.at(row, col))
273            .base2k(self.base2k())
274            .k(self.k())
275            .build()
276            .unwrap()
277    }
278}
279
280impl<D: DataMut> GGLWECiphertext<D> {
281    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
282        GLWECiphertext::builder()
283            .base2k(self.base2k())
284            .k(self.k())
285            .data(self.data.at_mut(row, col))
286            .build()
287            .unwrap()
288    }
289}
290
291impl GGLWECiphertext<Vec<u8>> {
292    pub fn alloc<A>(infos: &A) -> Self
293    where
294        A: GGLWEInfos,
295    {
296        Self::alloc_with(
297            infos.n(),
298            infos.base2k(),
299            infos.k(),
300            infos.rank_in(),
301            infos.rank_out(),
302            infos.dnum(),
303            infos.dsize(),
304        )
305    }
306
307    pub fn alloc_with(
308        n: Degree,
309        base2k: Base2K,
310        k: TorusPrecision,
311        rank_in: Rank,
312        rank_out: Rank,
313        dnum: Dnum,
314        dsize: Dsize,
315    ) -> Self {
316        let size: usize = k.0.div_ceil(base2k.0) as usize;
317        debug_assert!(
318            size as u32 > dsize.0,
319            "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
320            dsize.0
321        );
322
323        assert!(
324            dnum.0 * dsize.0 <= size as u32,
325            "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
326            dnum.0,
327            dsize.0,
328        );
329
330        Self {
331            data: MatZnx::alloc(
332                n.into(),
333                dnum.into(),
334                rank_in.into(),
335                (rank_out + 1).into(),
336                k.0.div_ceil(base2k.0) as usize,
337            ),
338            k,
339            base2k,
340            dsize,
341        }
342    }
343
344    pub fn alloc_bytes<A>(infos: &A) -> usize
345    where
346        A: GGLWEInfos,
347    {
348        Self::alloc_bytes_with(
349            infos.n(),
350            infos.base2k(),
351            infos.k(),
352            infos.rank_in(),
353            infos.rank_out(),
354            infos.dnum(),
355            infos.dsize(),
356        )
357    }
358
359    pub fn alloc_bytes_with(
360        n: Degree,
361        base2k: Base2K,
362        k: TorusPrecision,
363        rank_in: Rank,
364        rank_out: Rank,
365        dnum: Dnum,
366        dsize: Dsize,
367    ) -> usize {
368        let size: usize = k.0.div_ceil(base2k.0) as usize;
369        debug_assert!(
370            size as u32 > dsize.0,
371            "invalid gglwe: ceil(k/base2k): {size} <= dsize: {}",
372            dsize.0
373        );
374
375        assert!(
376            dnum.0 * dsize.0 <= size as u32,
377            "invalid gglwe: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
378            dnum.0,
379            dsize.0,
380        );
381
382        MatZnx::alloc_bytes(
383            n.into(),
384            dnum.into(),
385            rank_in.into(),
386            (rank_out + 1).into(),
387            k.0.div_ceil(base2k.0) as usize,
388        )
389    }
390}
391
392impl<D: DataMut> ReaderFrom for GGLWECiphertext<D> {
393    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
394        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
395        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
396        self.dsize = Dsize(reader.read_u32::<LittleEndian>()?);
397        self.data.read_from(reader)
398    }
399}
400
401impl<D: DataRef> WriterTo for GGLWECiphertext<D> {
402    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
403        writer.write_u32::<LittleEndian>(self.k.0)?;
404        writer.write_u32::<LittleEndian>(self.base2k.0)?;
405        writer.write_u32::<LittleEndian>(self.dsize.0)?;
406        self.data.write_to(writer)
407    }
408}