1use poulpy_hal::{
  2    api::{VecZnxCopy, VecZnxFillUniform},
  3    layouts::{
  4        Backend, Data, DataMut, DataRef, FillUniform, MatZnx, MatZnxToMut, MatZnxToRef, Module, ReaderFrom, WriterTo, ZnxInfos,
  5    },
  6    source::Source,
  7};
  8
  9use crate::layouts::{
 10    Base2K, Degree, Dnum, Dsize, GGSW, GGSWInfos, GLWEInfos, LWEInfos, Rank, TorusPrecision,
 11    compressed::{Decompress, GLWECiphertextCompressed},
 12};
 13use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 14use std::fmt;
 15
 16#[derive(PartialEq, Eq, Clone)]
 17pub struct GGSWCiphertextCompressed<D: Data> {
 18    pub(crate) data: MatZnx<D>,
 19    pub(crate) k: TorusPrecision,
 20    pub(crate) base2k: Base2K,
 21    pub(crate) dsize: Dsize,
 22    pub(crate) rank: Rank,
 23    pub(crate) seed: Vec<[u8; 32]>,
 24}
 25
 26impl<D: Data> LWEInfos for GGSWCiphertextCompressed<D> {
 27    fn n(&self) -> Degree {
 28        Degree(self.data.n() as u32)
 29    }
 30
 31    fn base2k(&self) -> Base2K {
 32        self.base2k
 33    }
 34
 35    fn k(&self) -> TorusPrecision {
 36        self.k
 37    }
 38    fn size(&self) -> usize {
 39        self.data.size()
 40    }
 41}
 42impl<D: Data> GLWEInfos for GGSWCiphertextCompressed<D> {
 43    fn rank(&self) -> Rank {
 44        self.rank
 45    }
 46}
 47
 48impl<D: Data> GGSWInfos for GGSWCiphertextCompressed<D> {
 49    fn dsize(&self) -> Dsize {
 50        self.dsize
 51    }
 52
 53    fn dnum(&self) -> Dnum {
 54        Dnum(self.data.rows() as u32)
 55    }
 56}
 57
 58impl<D: DataRef> fmt::Debug for GGSWCiphertextCompressed<D> {
 59    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 60        write!(f, "{}", self.data)
 61    }
 62}
 63
 64impl<D: DataRef> fmt::Display for GGSWCiphertextCompressed<D> {
 65    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 66        write!(
 67            f,
 68            "(GGSWCiphertextCompressed: base2k={} k={} dsize={}) {}",
 69            self.base2k, self.k, self.dsize, self.data
 70        )
 71    }
 72}
 73
 74impl<D: DataMut> FillUniform for GGSWCiphertextCompressed<D> {
 75    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
 76        self.data.fill_uniform(log_bound, source);
 77    }
 78}
 79
 80impl GGSWCiphertextCompressed<Vec<u8>> {
 81    pub fn alloc<A>(infos: &A) -> Self
 82    where
 83        A: GGSWInfos,
 84    {
 85        Self::alloc_with(
 86            infos.n(),
 87            infos.base2k(),
 88            infos.k(),
 89            infos.rank(),
 90            infos.dnum(),
 91            infos.dsize(),
 92        )
 93    }
 94
 95    pub fn alloc_with(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self {
 96        let size: usize = k.0.div_ceil(base2k.0) as usize;
 97        debug_assert!(
 98            size as u32 > dsize.0,
 99            "invalid ggsw: ceil(k/base2k): {size} <= dsize: {}",
100            dsize.0
101        );
102
103        assert!(
104            dnum.0 * dsize.0 <= size as u32,
105            "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
106            dnum.0,
107            dsize.0,
108        );
109
110        Self {
111            data: MatZnx::alloc(
112                n.into(),
113                dnum.into(),
114                (rank + 1).into(),
115                1,
116                k.0.div_ceil(base2k.0) as usize,
117            ),
118            k,
119            base2k,
120            dsize,
121            rank,
122            seed: Vec::new(),
123        }
124    }
125
126    pub fn alloc_bytes<A>(infos: &A) -> usize
127    where
128        A: GGSWInfos,
129    {
130        Self::alloc_bytes_with(
131            infos.n(),
132            infos.base2k(),
133            infos.k(),
134            infos.rank(),
135            infos.dnum(),
136            infos.dsize(),
137        )
138    }
139
140    pub fn alloc_bytes_with(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
141        let size: usize = k.0.div_ceil(base2k.0) as usize;
142        debug_assert!(
143            size as u32 > dsize.0,
144            "invalid ggsw: ceil(k/base2k): {size} <= dsize: {}",
145            dsize.0
146        );
147
148        assert!(
149            dnum.0 * dsize.0 <= size as u32,
150            "invalid ggsw: dnum: {} * dsize:{} > ceil(k/base2k): {size}",
151            dnum.0,
152            dsize.0,
153        );
154
155        MatZnx::alloc_bytes(
156            n.into(),
157            dnum.into(),
158            (rank + 1).into(),
159            1,
160            k.0.div_ceil(base2k.0) as usize,
161        )
162    }
163}
164
165impl<D: DataRef> GGSWCiphertextCompressed<D> {
166    pub fn at(&self, row: usize, col: usize) -> GLWECiphertextCompressed<&[u8]> {
167        let rank: usize = self.rank().into();
168        GLWECiphertextCompressed {
169            data: self.data.at(row, col),
170            k: self.k,
171            base2k: self.base2k,
172            rank: self.rank,
173            seed: self.seed[row * (rank + 1) + col],
174        }
175    }
176}
177
178impl<D: DataMut> GGSWCiphertextCompressed<D> {
179    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertextCompressed<&mut [u8]> {
180        let rank: usize = self.rank().into();
181        GLWECiphertextCompressed {
182            data: self.data.at_mut(row, col),
183            k: self.k,
184            base2k: self.base2k,
185            rank: self.rank,
186            seed: self.seed[row * (rank + 1) + col],
187        }
188    }
189}
190
191impl<D: DataMut> ReaderFrom for GGSWCiphertextCompressed<D> {
192    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
193        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
194        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
195        self.dsize = Dsize(reader.read_u32::<LittleEndian>()?);
196        self.rank = Rank(reader.read_u32::<LittleEndian>()?);
197        let seed_len: usize = reader.read_u32::<LittleEndian>()? as usize;
198        self.seed = vec![[0u8; 32]; seed_len];
199        for s in &mut self.seed {
200            reader.read_exact(s)?;
201        }
202        self.data.read_from(reader)
203    }
204}
205
206impl<D: DataRef> WriterTo for GGSWCiphertextCompressed<D> {
207    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
208        writer.write_u32::<LittleEndian>(self.k.into())?;
209        writer.write_u32::<LittleEndian>(self.base2k.into())?;
210        writer.write_u32::<LittleEndian>(self.dsize.into())?;
211        writer.write_u32::<LittleEndian>(self.rank.into())?;
212        writer.write_u32::<LittleEndian>(self.seed.len() as u32)?;
213        for s in &self.seed {
214            writer.write_all(s)?;
215        }
216        self.data.write_to(writer)
217    }
218}
219
220impl<D: DataMut, B: Backend, DR: DataRef> Decompress<B, GGSWCiphertextCompressed<DR>> for GGSW<D>
221where
222    Module<B>: VecZnxFillUniform + VecZnxCopy,
223{
224    fn decompress(&mut self, module: &Module<B>, other: &GGSWCiphertextCompressed<DR>) {
225        #[cfg(debug_assertions)]
226        {
227            assert_eq!(self.rank(), other.rank())
228        }
229
230        let dnum: usize = self.dnum().into();
231        let rank: usize = self.rank().into();
232        (0..dnum).for_each(|row_i| {
233            (0..rank + 1).for_each(|col_j| {
234                self.at_mut(row_i, col_j)
235                    .decompress(module, &other.at(row_i, col_j));
236            });
237        });
238    }
239}
240
241pub trait GGSWCiphertextCompressedToMut {
242    fn to_mut(&mut self) -> GGSWCiphertextCompressed<&mut [u8]>;
243}
244
245impl<D: DataMut> GGSWCiphertextCompressedToMut for GGSWCiphertextCompressed<D> {
246    fn to_mut(&mut self) -> GGSWCiphertextCompressed<&mut [u8]> {
247        GGSWCiphertextCompressed {
248            k: self.k(),
249            base2k: self.base2k(),
250            dsize: self.dsize(),
251            rank: self.rank(),
252            seed: self.seed.clone(),
253            data: self.data.to_mut(),
254        }
255    }
256}
257
258pub trait GGSWCiphertextCompressedToRef {
259    fn to_ref(&self) -> GGSWCiphertextCompressed<&[u8]>;
260}
261
262impl<D: DataRef> GGSWCiphertextCompressedToRef for GGSWCiphertextCompressed<D> {
263    fn to_ref(&self) -> GGSWCiphertextCompressed<&[u8]> {
264        GGSWCiphertextCompressed {
265            k: self.k(),
266            base2k: self.base2k(),
267            dsize: self.dsize(),
268            rank: self.rank(),
269            seed: self.seed.clone(),
270            data: self.data.to_ref(),
271        }
272    }
273}