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