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
 79impl<BE: Backend> GGSWCompressedAlloc for Module<BE> where Self: GetDegree {}
 80
 81pub trait GGSWCompressedAlloc
 82where
 83    Self: GetDegree,
 84{
 85    fn alloc_ggsw_compressed(
 86        &self,
 87        base2k: Base2K,
 88        k: TorusPrecision,
 89        rank: Rank,
 90        dnum: Dnum,
 91        dsize: Dsize,
 92    ) -> GGSWCompressed<Vec<u8>> {
 93        let size: usize = k.0.div_ceil(base2k.0) as usize;
 94        assert!(
 95            size as u32 > dsize.0,
 96            "invalid ggsw: ceil(k/base2k): {size} <= dsize: {}",
 97            dsize.0
 98        );
 99
100        assert!(
101            dnum.0 * dsize.0 <= size as u32,
102            "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
103            dnum.0,
104            dsize.0,
105        );
106
107        GGSWCompressed {
108            data: MatZnx::alloc(
109                self.ring_degree().into(),
110                dnum.into(),
111                (rank + 1).into(),
112                1,
113                k.0.div_ceil(base2k.0) as usize,
114            ),
115            k,
116            base2k,
117            dsize,
118            rank,
119            seed: Vec::new(),
120        }
121    }
122
123    fn alloc_ggsw_compressed_from_infos<A>(&self, infos: &A) -> GGSWCompressed<Vec<u8>>
124    where
125        A: GGSWInfos,
126    {
127        self.alloc_ggsw_compressed(
128            infos.base2k(),
129            infos.k(),
130            infos.rank(),
131            infos.dnum(),
132            infos.dsize(),
133        )
134    }
135
136    fn bytes_of_ggsw_compressed(&self, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
137        let size: usize = k.0.div_ceil(base2k.0) as usize;
138        assert!(
139            size as u32 > dsize.0,
140            "invalid ggsw: ceil(k/base2k): {size} <= dsize: {}",
141            dsize.0
142        );
143
144        assert!(
145            dnum.0 * dsize.0 <= size as u32,
146            "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
147            dnum.0,
148            dsize.0,
149        );
150
151        MatZnx::bytes_of(
152            self.ring_degree().into(),
153            dnum.into(),
154            (rank + 1).into(),
155            1,
156            k.0.div_ceil(base2k.0) as usize,
157        )
158    }
159
160    fn bytes_of_ggsw_compressed_key_from_infos<A>(&self, infos: &A) -> usize
161    where
162        A: GGSWInfos,
163    {
164        self.bytes_of_ggsw_compressed(
165            infos.base2k(),
166            infos.k(),
167            infos.rank(),
168            infos.dnum(),
169            infos.dsize(),
170        )
171    }
172}
173
174impl GGSWCompressed<Vec<u8>> {
175    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
176    where
177        A: GGSWInfos,
178        M: GGSWCompressedAlloc,
179    {
180        module.alloc_ggsw_compressed_from_infos(infos)
181    }
182
183    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self
184    where
185        M: GGSWCompressedAlloc,
186    {
187        module.alloc_ggsw_compressed(base2k, k, rank, dnum, dsize)
188    }
189
190    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
191    where
192        A: GGSWInfos,
193        M: GGSWCompressedAlloc,
194    {
195        module.bytes_of_ggsw_compressed_key_from_infos(infos)
196    }
197
198    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize
199    where
200        M: GGSWCompressedAlloc,
201    {
202        module.bytes_of_ggsw_compressed(base2k, k, rank, dnum, dsize)
203    }
204}
205
206impl<D: DataRef> GGSWCompressed<D> {
207    pub fn at(&self, row: usize, col: usize) -> GLWECompressed<&[u8]> {
208        let rank: usize = self.rank().into();
209        GLWECompressed {
210            data: self.data.at(row, col),
211            k: self.k,
212            base2k: self.base2k,
213            rank: self.rank,
214            seed: self.seed[row * (rank + 1) + col],
215        }
216    }
217}
218
219impl<D: DataMut> GGSWCompressed<D> {
220    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECompressed<&mut [u8]> {
221        let rank: usize = self.rank().into();
222        GLWECompressed {
223            data: self.data.at_mut(row, col),
224            k: self.k,
225            base2k: self.base2k,
226            rank: self.rank,
227            seed: self.seed[row * (rank + 1) + col],
228        }
229    }
230}
231
232impl<D: DataMut> ReaderFrom for GGSWCompressed<D> {
233    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
234        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
235        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
236        self.dsize = Dsize(reader.read_u32::<LittleEndian>()?);
237        self.rank = Rank(reader.read_u32::<LittleEndian>()?);
238        let seed_len: usize = reader.read_u32::<LittleEndian>()? as usize;
239        self.seed = vec![[0u8; 32]; seed_len];
240        for s in &mut self.seed {
241            reader.read_exact(s)?;
242        }
243        self.data.read_from(reader)
244    }
245}
246
247impl<D: DataRef> WriterTo for GGSWCompressed<D> {
248    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
249        writer.write_u32::<LittleEndian>(self.k.into())?;
250        writer.write_u32::<LittleEndian>(self.base2k.into())?;
251        writer.write_u32::<LittleEndian>(self.dsize.into())?;
252        writer.write_u32::<LittleEndian>(self.rank.into())?;
253        writer.write_u32::<LittleEndian>(self.seed.len() as u32)?;
254        for s in &self.seed {
255            writer.write_all(s)?;
256        }
257        self.data.write_to(writer)
258    }
259}
260
261pub trait GGSWDecompress
262where
263    Self: GLWEDecompress,
264{
265    fn decompress_ggsw<R, O>(&self, res: &mut R, other: &O)
266    where
267        R: GGSWToMut,
268        O: GGSWCompressedToRef,
269    {
270        let res: &mut GGSW<&mut [u8]> = &mut res.to_mut();
271        let other: &GGSWCompressed<&[u8]> = &other.to_ref();
272
273        assert_eq!(res.rank(), other.rank());
274        let dnum: usize = res.dnum().into();
275        let rank: usize = res.rank().into();
276
277        for row_i in 0..dnum {
278            for col_j in 0..rank + 1 {
279                self.decompress_glwe(&mut res.at_mut(row_i, col_j), &other.at(row_i, col_j));
280            }
281        }
282    }
283}
284
285impl<B: Backend> GGSWDecompress for Module<B> where Self: GGSWDecompress {}
286
287impl<D: DataMut> GGSW<D> {
288    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
289    where
290        O: GGSWCompressedToRef,
291        M: GGSWDecompress,
292    {
293        module.decompress_ggsw(self, other);
294    }
295}
296
297pub trait GGSWCompressedToMut {
298    fn to_mut(&mut self) -> GGSWCompressed<&mut [u8]>;
299}
300
301impl<D: DataMut> GGSWCompressedToMut for GGSWCompressed<D> {
302    fn to_mut(&mut self) -> GGSWCompressed<&mut [u8]> {
303        GGSWCompressed {
304            k: self.k(),
305            base2k: self.base2k(),
306            dsize: self.dsize(),
307            rank: self.rank(),
308            seed: self.seed.clone(),
309            data: self.data.to_mut(),
310        }
311    }
312}
313
314pub trait GGSWCompressedToRef {
315    fn to_ref(&self) -> GGSWCompressed<&[u8]>;
316}
317
318impl<D: DataRef> GGSWCompressedToRef for GGSWCompressed<D> {
319    fn to_ref(&self) -> GGSWCompressed<&[u8]> {
320        GGSWCompressed {
321            k: self.k(),
322            base2k: self.base2k(),
323            dsize: self.dsize(),
324            rank: self.rank(),
325            seed: self.seed.clone(),
326            data: self.data.to_ref(),
327        }
328    }
329}