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