1use backend::hal::{
  2    api::{FillUniform, Reset, VecZnxCopy, VecZnxFillUniform, ZnxInfos},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, ReaderFrom, VecZnx, VecZnxToMut, VecZnxToRef, WriterTo},
  4};
  5use sampling::source::Source;
  6
  7use crate::{Decompress, GLWEOps, Infos, SetMetaData};
  8use std::fmt;
  9
 10#[derive(PartialEq, Eq, Clone)]
 11pub struct GLWECiphertext<D: Data> {
 12    pub data: VecZnx<D>,
 13    pub basek: usize,
 14    pub k: usize,
 15}
 16
 17impl<D: DataRef> fmt::Debug for GLWECiphertext<D> {
 18    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 19        write!(f, "{}", self)
 20    }
 21}
 22
 23impl<D: DataRef> fmt::Display for GLWECiphertext<D> {
 24    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 25        write!(
 26            f,
 27            "GLWECiphertext: basek={} k={}: {}",
 28            self.basek(),
 29            self.k(),
 30            self.data
 31        )
 32    }
 33}
 34
 35impl<D: DataMut> Reset for GLWECiphertext<D>
 36where
 37    VecZnx<D>: Reset,
 38{
 39    fn reset(&mut self) {
 40        self.data.reset();
 41        self.basek = 0;
 42        self.k = 0;
 43    }
 44}
 45
 46impl<D: DataMut> FillUniform for GLWECiphertext<D>
 47where
 48    VecZnx<D>: FillUniform,
 49{
 50    fn fill_uniform(&mut self, source: &mut Source) {
 51        self.data.fill_uniform(source);
 52    }
 53}
 54
 55impl GLWECiphertext<Vec<u8>> {
 56    pub fn alloc(n: usize, basek: usize, k: usize, rank: usize) -> Self {
 57        Self {
 58            data: VecZnx::alloc(n, rank + 1, k.div_ceil(basek)),
 59            basek,
 60            k,
 61        }
 62    }
 63
 64    pub fn bytes_of(n: usize, basek: usize, k: usize, rank: usize) -> usize {
 65        VecZnx::alloc_bytes(n, rank + 1, k.div_ceil(basek))
 66    }
 67}
 68
 69impl<D: Data> Infos for GLWECiphertext<D> {
 70    type Inner = VecZnx<D>;
 71
 72    fn inner(&self) -> &Self::Inner {
 73        &self.data
 74    }
 75
 76    fn basek(&self) -> usize {
 77        self.basek
 78    }
 79
 80    fn k(&self) -> usize {
 81        self.k
 82    }
 83}
 84
 85impl<D: Data> GLWECiphertext<D> {
 86    pub fn rank(&self) -> usize {
 87        self.cols() - 1
 88    }
 89}
 90
 91impl<D: DataRef> GLWECiphertext<D> {
 92    pub fn clone(&self) -> GLWECiphertext<Vec<u8>> {
 93        GLWECiphertext {
 94            data: self.data.clone(),
 95            basek: self.basek(),
 96            k: self.k(),
 97        }
 98    }
 99}
100
101impl<D: DataMut> SetMetaData for GLWECiphertext<D> {
102    fn set_k(&mut self, k: usize) {
103        self.k = k
104    }
105
106    fn set_basek(&mut self, basek: usize) {
107        self.basek = basek
108    }
109}
110
111pub trait GLWECiphertextToRef: Infos {
112    fn to_ref(&self) -> GLWECiphertext<&[u8]>;
113}
114
115impl<D: DataRef> GLWECiphertextToRef for GLWECiphertext<D> {
116    fn to_ref(&self) -> GLWECiphertext<&[u8]> {
117        GLWECiphertext {
118            data: self.data.to_ref(),
119            basek: self.basek,
120            k: self.k,
121        }
122    }
123}
124
125pub trait GLWECiphertextToMut: Infos {
126    fn to_mut(&mut self) -> GLWECiphertext<&mut [u8]>;
127}
128
129impl<D: DataMut> GLWECiphertextToMut for GLWECiphertext<D> {
130    fn to_mut(&mut self) -> GLWECiphertext<&mut [u8]> {
131        GLWECiphertext {
132            data: self.data.to_mut(),
133            basek: self.basek,
134            k: self.k,
135        }
136    }
137}
138
139impl<D: DataMut> GLWEOps for GLWECiphertext<D> where GLWECiphertext<D>: GLWECiphertextToMut + Infos + SetMetaData {}
140
141use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
142
143impl<D: DataMut> ReaderFrom for GLWECiphertext<D> {
144    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
145        self.k = reader.read_u64::<LittleEndian>()? as usize;
146        self.basek = reader.read_u64::<LittleEndian>()? as usize;
147        self.data.read_from(reader)
148    }
149}
150
151impl<D: DataRef> WriterTo for GLWECiphertext<D> {
152    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
153        writer.write_u64::<LittleEndian>(self.k as u64)?;
154        writer.write_u64::<LittleEndian>(self.basek as u64)?;
155        self.data.write_to(writer)
156    }
157}
158
159#[derive(PartialEq, Eq, Clone)]
160pub struct GLWECiphertextCompressed<D: Data> {
161    pub(crate) data: VecZnx<D>,
162    pub(crate) basek: usize,
163    pub(crate) k: usize,
164    pub(crate) rank: usize,
165    pub(crate) seed: [u8; 32],
166}
167
168impl<D: DataRef> fmt::Debug for GLWECiphertextCompressed<D> {
169    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
170        write!(f, "{}", self)
171    }
172}
173
174impl<D: DataRef> fmt::Display for GLWECiphertextCompressed<D> {
175    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
176        write!(
177            f,
178            "GLWECiphertextCompressed: basek={} k={} rank={} seed={:?}: {}",
179            self.basek(),
180            self.k(),
181            self.rank,
182            self.seed,
183            self.data
184        )
185    }
186}
187
188impl<D: DataMut> Reset for GLWECiphertextCompressed<D>
189where
190    VecZnx<D>: Reset,
191{
192    fn reset(&mut self) {
193        self.data.reset();
194        self.basek = 0;
195        self.k = 0;
196        self.rank = 0;
197        self.seed = [0u8; 32];
198    }
199}
200
201impl<D: DataMut> FillUniform for GLWECiphertextCompressed<D>
202where
203    VecZnx<D>: FillUniform,
204{
205    fn fill_uniform(&mut self, source: &mut Source) {
206        self.data.fill_uniform(source);
207    }
208}
209
210impl<D: Data> Infos for GLWECiphertextCompressed<D> {
211    type Inner = VecZnx<D>;
212
213    fn inner(&self) -> &Self::Inner {
214        &self.data
215    }
216
217    fn basek(&self) -> usize {
218        self.basek
219    }
220
221    fn k(&self) -> usize {
222        self.k
223    }
224}
225
226impl<D: Data> GLWECiphertextCompressed<D> {
227    pub fn rank(&self) -> usize {
228        self.rank
229    }
230}
231
232impl GLWECiphertextCompressed<Vec<u8>> {
233    pub fn alloc(n: usize, basek: usize, k: usize, rank: usize) -> Self {
234        Self {
235            data: VecZnx::alloc(n, 1, k.div_ceil(basek)),
236            basek,
237            k,
238            rank,
239            seed: [0u8; 32],
240        }
241    }
242
243    pub fn bytes_of(n: usize, basek: usize, k: usize) -> usize {
244        GLWECiphertext::bytes_of(n, basek, k, 1)
245    }
246}
247
248impl<D: DataMut> ReaderFrom for GLWECiphertextCompressed<D> {
249    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
250        self.k = reader.read_u64::<LittleEndian>()? as usize;
251        self.basek = reader.read_u64::<LittleEndian>()? as usize;
252        self.rank = reader.read_u64::<LittleEndian>()? as usize;
253        reader.read(&mut self.seed)?;
254        self.data.read_from(reader)
255    }
256}
257
258impl<D: DataRef> WriterTo for GLWECiphertextCompressed<D> {
259    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
260        writer.write_u64::<LittleEndian>(self.k as u64)?;
261        writer.write_u64::<LittleEndian>(self.basek as u64)?;
262        writer.write_u64::<LittleEndian>(self.rank as u64)?;
263        writer.write_all(&self.seed)?;
264        self.data.write_to(writer)
265    }
266}
267
268impl<D: DataMut, B: Backend, DR: DataRef> Decompress<B, GLWECiphertextCompressed<DR>> for GLWECiphertext<D> {
269    fn decompress(&mut self, module: &Module<B>, other: &GLWECiphertextCompressed<DR>)
270    where
271        Module<B>: VecZnxFillUniform + VecZnxCopy,
272    {
273        #[cfg(debug_assertions)]
274        {
275            assert_eq!(
276                self.n(),
277                other.data.n(),
278                "invalid receiver: self.n()={} != other.n()={}",
279                self.n(),
280                other.data.n()
281            );
282            assert_eq!(
283                self.size(),
284                other.size(),
285                "invalid receiver: self.size()={} != other.size()={}",
286                self.size(),
287                other.size()
288            );
289            assert_eq!(
290                self.rank(),
291                other.rank(),
292                "invalid receiver: self.rank()={} != other.rank()={}",
293                self.rank(),
294                other.rank()
295            );
296            let mut source: Source = Source::new(other.seed);
297            self.decompress_internal(module, other, &mut source);
298        }
299    }
300}
301
302impl<D: DataMut> GLWECiphertext<D> {
303    pub(crate) fn decompress_internal<DataOther, B: Backend>(
304        &mut self,
305        module: &Module<B>,
306        other: &GLWECiphertextCompressed<DataOther>,
307        source: &mut Source,
308    ) where
309        DataOther: DataRef,
310        Module<B>: VecZnxFillUniform + VecZnxCopy,
311    {
312        #[cfg(debug_assertions)]
313        {
314            assert_eq!(self.rank(), other.rank())
315        }
316
317        let k: usize = other.k;
318        let basek: usize = other.basek;
319        let cols: usize = other.rank() + 1;
320        module.vec_znx_copy(&mut self.data, 0, &other.data, 0);
321        (1..cols).for_each(|i| {
322            module.vec_znx_fill_uniform(basek, &mut self.data, i, k, source);
323        });
324
325        self.basek = basek;
326        self.k = k;
327    }
328}