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>(
197        module: &M,
198        base2k: Base2K,
199        k: TorusPrecision,
200        rank: Rank,
201        dnum: Dnum,
202        dsize: Dsize,
203    ) -> usize
204    where
205        M: GGSWCompressedAlloc,
206    {
207        module.bytes_of_ggsw_compressed(base2k, k, rank, dnum, dsize)
208    }
209}
210
211impl<D: DataRef> GGSWCompressed<D> {
212    pub fn at(&self, row: usize, col: usize) -> GLWECompressed<&[u8]> {
213        let rank: usize = self.rank().into();
214        GLWECompressed {
215            data: self.data.at(row, col),
216            k: self.k,
217            base2k: self.base2k,
218            rank: self.rank,
219            seed: self.seed[row * (rank + 1) + col],
220        }
221    }
222}
223
224impl<D: DataMut> GGSWCompressed<D> {
225    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECompressed<&mut [u8]> {
226        let rank: usize = self.rank().into();
227        GLWECompressed {
228            data: self.data.at_mut(row, col),
229            k: self.k,
230            base2k: self.base2k,
231            rank: self.rank,
232            seed: self.seed[row * (rank + 1) + col],
233        }
234    }
235}
236
237impl<D: DataMut> ReaderFrom for GGSWCompressed<D> {
238    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
239        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
240        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
241        self.dsize = Dsize(reader.read_u32::<LittleEndian>()?);
242        self.rank = Rank(reader.read_u32::<LittleEndian>()?);
243        let seed_len: usize = reader.read_u32::<LittleEndian>()? as usize;
244        self.seed = vec![[0u8; 32]; seed_len];
245        for s in &mut self.seed {
246            reader.read_exact(s)?;
247        }
248        self.data.read_from(reader)
249    }
250}
251
252impl<D: DataRef> WriterTo for GGSWCompressed<D> {
253    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
254        writer.write_u32::<LittleEndian>(self.k.into())?;
255        writer.write_u32::<LittleEndian>(self.base2k.into())?;
256        writer.write_u32::<LittleEndian>(self.dsize.into())?;
257        writer.write_u32::<LittleEndian>(self.rank.into())?;
258        writer.write_u32::<LittleEndian>(self.seed.len() as u32)?;
259        for s in &self.seed {
260            writer.write_all(s)?;
261        }
262        self.data.write_to(writer)
263    }
264}
265
266pub trait GGSWDecompress
267where
268    Self: GLWEDecompress,
269{
270    fn decompress_ggsw<R, O>(&self, res: &mut R, other: &O)
271    where
272        R: GGSWToMut,
273        O: GGSWCompressedToRef,
274    {
275        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
276        let other: &GGSWCompressed<&[u8]> = &other.to_ref();
277
278        assert_eq!(res.rank(), other.rank());
279        let dnum: usize = res.dnum().into();
280        let rank: usize = res.rank().into();
281
282        for row_i in 0..dnum {
283            for col_j in 0..rank + 1 {
284                self.decompress_glwe(&mut res.at_mut(row_i, col_j), &other.at(row_i, col_j));
285            }
286        }
287    }
288}
289
290impl<B: Backend> GGSWDecompress for Module<B> where Self: GGSWDecompress {}
291
292impl<D: DataMut> GGSW<D> {
293    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
294    where
295        O: GGSWCompressedToRef,
296        M: GGSWDecompress,
297    {
298        module.decompress_ggsw(self, other);
299    }
300}
301
302pub trait GGSWCompressedToMut {
303    fn to_mut(&mut self) -> GGSWCompressed<&mut [u8]>;
304}
305
306impl<D: DataMut> GGSWCompressedToMut for GGSWCompressed<D> {
307    fn to_mut(&mut self) -> GGSWCompressed<&mut [u8]> {
308        GGSWCompressed {
309            k: self.k(),
310            base2k: self.base2k(),
311            dsize: self.dsize(),
312            rank: self.rank(),
313            seed: self.seed.clone(),
314            data: self.data.to_mut(),
315        }
316    }
317}
318
319pub trait GGSWCompressedToRef {
320    fn to_ref(&self) -> GGSWCompressed<&[u8]>;
321}
322
323impl<D: DataRef> GGSWCompressedToRef for GGSWCompressed<D> {
324    fn to_ref(&self) -> GGSWCompressed<&[u8]> {
325        GGSWCompressed {
326            k: self.k(),
327            base2k: self.base2k(),
328            dsize: self.dsize(),
329            rank: self.rank(),
330            seed: self.seed.clone(),
331            data: self.data.to_ref(),
332        }
333    }
334}