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