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::{
  9    Base2K, Degree, Dnum, Dsize, GGSW, GGSWInfos, GGSWToMut, GLWEInfos, GetDegree, LWEInfos, Rank, TorusPrecision,
 10    compressed::{GLWECompressed, GLWEDecompress},
 11};
 12use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 13use std::fmt;
 14
 15#[derive(PartialEq, Eq, Clone)]
 16pub struct GGSWCompressed<D: Data> {
 17    pub(crate) data: MatZnx<D>,
 18    pub(crate) k: TorusPrecision,
 19    pub(crate) base2k: Base2K,
 20    pub(crate) dsize: Dsize,
 21    pub(crate) rank: Rank,
 22    pub(crate) seed: Vec<[u8; 32]>,
 23}
 24
 25impl<D: Data> LWEInfos for GGSWCompressed<D> {
 26    fn n(&self) -> Degree {
 27        Degree(self.data.n() as u32)
 28    }
 29
 30    fn base2k(&self) -> Base2K {
 31        self.base2k
 32    }
 33
 34    fn k(&self) -> TorusPrecision {
 35        self.k
 36    }
 37    fn size(&self) -> usize {
 38        self.data.size()
 39    }
 40}
 41impl<D: Data> GLWEInfos for GGSWCompressed<D> {
 42    fn rank(&self) -> Rank {
 43        self.rank
 44    }
 45}
 46
 47impl<D: Data> GGSWInfos for GGSWCompressed<D> {
 48    fn dsize(&self) -> Dsize {
 49        self.dsize
 50    }
 51
 52    fn dnum(&self) -> Dnum {
 53        Dnum(self.data.rows() as u32)
 54    }
 55}
 56
 57impl<D: DataRef> fmt::Debug for GGSWCompressed<D> {
 58    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 59        write!(f, "{}", self.data)
 60    }
 61}
 62
 63impl<D: DataRef> fmt::Display for GGSWCompressed<D> {
 64    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 65        write!(
 66            f,
 67            "(GGSWCompressed: base2k={} k={} dsize={}) {}",
 68            self.base2k, self.k, self.dsize, self.data
 69        )
 70    }
 71}
 72
 73impl<D: DataMut> FillUniform for GGSWCompressed<D> {
 74    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
 75        self.data.fill_uniform(log_bound, source);
 76    }
 77}
 78
 79pub trait GGSWCompressedAlloc
 80where
 81    Self: GetDegree,
 82{
 83    fn alloc_ggsw_compressed(
 84        &self,
 85        base2k: Base2K,
 86        k: TorusPrecision,
 87        rank: Rank,
 88        dnum: Dnum,
 89        dsize: Dsize,
 90    ) -> GGSWCompressed<Vec<u8>> {
 91        let size: usize = k.0.div_ceil(base2k.0) as usize;
 92        assert!(
 93            size as u32 > dsize.0,
 94            "invalid ggsw: ceil(k/base2k): {size} <= dsize: {}",
 95            dsize.0
 96        );
 97
 98        assert!(
 99            dnum.0 * dsize.0 <= size as u32,
100            "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
101            dnum.0,
102            dsize.0,
103        );
104
105        GGSWCompressed {
106            data: MatZnx::alloc(
107                self.n().into(),
108                dnum.into(),
109                (rank + 1).into(),
110                1,
111                k.0.div_ceil(base2k.0) as usize,
112            ),
113            k,
114            base2k,
115            dsize,
116            rank,
117            seed: Vec::new(),
118        }
119    }
120
121    fn alloc_ggsw_compressed_from_infos<A>(&self, infos: &A) -> GGSWCompressed<Vec<u8>>
122    where
123        A: GGSWInfos,
124    {
125        self.alloc_ggsw_compressed(
126            infos.base2k(),
127            infos.k(),
128            infos.rank(),
129            infos.dnum(),
130            infos.dsize(),
131        )
132    }
133
134    fn bytes_of_ggsw_compressed(&self, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
135        let size: usize = k.0.div_ceil(base2k.0) as usize;
136        assert!(
137            size as u32 > dsize.0,
138            "invalid ggsw: ceil(k/base2k): {size} <= dsize: {}",
139            dsize.0
140        );
141
142        assert!(
143            dnum.0 * dsize.0 <= size as u32,
144            "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
145            dnum.0,
146            dsize.0,
147        );
148
149        MatZnx::bytes_of(
150            self.n().into(),
151            dnum.into(),
152            (rank + 1).into(),
153            1,
154            k.0.div_ceil(base2k.0) as usize,
155        )
156    }
157
158    fn bytes_of_ggsw_compressed_key_from_infos<A>(&self, infos: &A) -> usize
159    where
160        A: GGSWInfos,
161    {
162        self.bytes_of_ggsw_compressed(
163            infos.base2k(),
164            infos.k(),
165            infos.rank(),
166            infos.dnum(),
167            infos.dsize(),
168        )
169    }
170}
171
172impl GGSWCompressed<Vec<u8>> {
173    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
174    where
175        A: GGSWInfos,
176        M: GGSWCompressedAlloc,
177    {
178        module.alloc_ggsw_compressed_from_infos(infos)
179    }
180
181    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self
182    where
183        M: GGSWCompressedAlloc,
184    {
185        module.alloc_ggsw_compressed(base2k, k, rank, dnum, dsize)
186    }
187
188    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
189    where
190        A: GGSWInfos,
191        M: GGSWCompressedAlloc,
192    {
193        module.bytes_of_ggsw_compressed_key_from_infos(infos)
194    }
195
196    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize
197    where
198        M: GGSWCompressedAlloc,
199    {
200        module.bytes_of_ggsw_compressed(base2k, k, rank, dnum, dsize)
201    }
202}
203
204impl<D: DataRef> GGSWCompressed<D> {
205    pub fn at(&self, row: usize, col: usize) -> GLWECompressed<&[u8]> {
206        let rank: usize = self.rank().into();
207        GLWECompressed {
208            data: self.data.at(row, col),
209            k: self.k,
210            base2k: self.base2k,
211            rank: self.rank,
212            seed: self.seed[row * (rank + 1) + col],
213        }
214    }
215}
216
217impl<D: DataMut> GGSWCompressed<D> {
218    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECompressed<&mut [u8]> {
219        let rank: usize = self.rank().into();
220        GLWECompressed {
221            data: self.data.at_mut(row, col),
222            k: self.k,
223            base2k: self.base2k,
224            rank: self.rank,
225            seed: self.seed[row * (rank + 1) + col],
226        }
227    }
228}
229
230impl<D: DataMut> ReaderFrom for GGSWCompressed<D> {
231    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
232        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
233        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
234        self.dsize = Dsize(reader.read_u32::<LittleEndian>()?);
235        self.rank = Rank(reader.read_u32::<LittleEndian>()?);
236        let seed_len: usize = reader.read_u32::<LittleEndian>()? as usize;
237        self.seed = vec![[0u8; 32]; seed_len];
238        for s in &mut self.seed {
239            reader.read_exact(s)?;
240        }
241        self.data.read_from(reader)
242    }
243}
244
245impl<D: DataRef> WriterTo for GGSWCompressed<D> {
246    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
247        writer.write_u32::<LittleEndian>(self.k.into())?;
248        writer.write_u32::<LittleEndian>(self.base2k.into())?;
249        writer.write_u32::<LittleEndian>(self.dsize.into())?;
250        writer.write_u32::<LittleEndian>(self.rank.into())?;
251        writer.write_u32::<LittleEndian>(self.seed.len() as u32)?;
252        for s in &self.seed {
253            writer.write_all(s)?;
254        }
255        self.data.write_to(writer)
256    }
257}
258
259pub trait GGSWDecompress
260where
261    Self: GLWEDecompress,
262{
263    fn decompress_ggsw<R, O>(&self, res: &mut R, other: &O)
264    where
265        R: GGSWToMut,
266        O: GGSWCompressedToRef,
267    {
268        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
269        let other: &GGSWCompressed<&[u8]> = &other.to_ref();
270
271        assert_eq!(res.rank(), other.rank());
272        let dnum: usize = res.dnum().into();
273        let rank: usize = res.rank().into();
274
275        for row_i in 0..dnum {
276            for col_j in 0..rank + 1 {
277                self.decompress_glwe(&mut res.at_mut(row_i, col_j), &other.at(row_i, col_j));
278            }
279        }
280    }
281}
282
283impl<B: Backend> GGSWDecompress for Module<B> where Self: GGSWDecompress {}
284
285impl<D: DataMut> GGSW<D> {
286    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
287    where
288        O: GGSWCompressedToRef,
289        M: GGSWDecompress,
290    {
291        module.decompress_ggsw(self, other);
292    }
293}
294
295pub trait GGSWCompressedToMut {
296    fn to_mut(&mut self) -> GGSWCompressed<&mut [u8]>;
297}
298
299impl<D: DataMut> GGSWCompressedToMut for GGSWCompressed<D> {
300    fn to_mut(&mut self) -> GGSWCompressed<&mut [u8]> {
301        GGSWCompressed {
302            k: self.k(),
303            base2k: self.base2k(),
304            dsize: self.dsize(),
305            rank: self.rank(),
306            seed: self.seed.clone(),
307            data: self.data.to_mut(),
308        }
309    }
310}
311
312pub trait GGSWCompressedToRef {
313    fn to_ref(&self) -> GGSWCompressed<&[u8]>;
314}
315
316impl<D: DataRef> GGSWCompressedToRef for GGSWCompressed<D> {
317    fn to_ref(&self) -> GGSWCompressed<&[u8]> {
318        GGSWCompressed {
319            k: self.k(),
320            base2k: self.base2k(),
321            dsize: self.dsize(),
322            rank: self.rank(),
323            seed: self.seed.clone(),
324            data: self.data.to_ref(),
325        }
326    }
327}