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