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