1use backend::hal::{
  2    api::{FillUniform, Reset, VecZnxCopy, VecZnxFillUniform},
  3    layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, WriterTo},
  4};
  5
  6use crate::{AutomorphismKey, Decompress, GGLWECiphertext, GLWECiphertextCompressed, GLWESwitchingKey, GLWETensorKey, Infos};
  7use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
  8use std::fmt;
  9
 10#[derive(PartialEq, Eq, Clone)]
 11pub struct GGLWECiphertextCompressed<D: Data> {
 12    pub(crate) data: MatZnx<D>,
 13    pub(crate) basek: usize,
 14    pub(crate) k: usize,
 15    pub(crate) rank_out: usize,
 16    pub(crate) digits: usize,
 17    pub(crate) seed: Vec<[u8; 32]>,
 18}
 19
 20impl<D: DataRef> fmt::Debug for GGLWECiphertextCompressed<D> {
 21    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 22        write!(f, "{}", self)
 23    }
 24}
 25
 26impl<D: DataMut> FillUniform for GGLWECiphertextCompressed<D> {
 27    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
 28        self.data.fill_uniform(source);
 29    }
 30}
 31
 32impl<D: DataMut> Reset for GGLWECiphertextCompressed<D>
 33where
 34    MatZnx<D>: Reset,
 35{
 36    fn reset(&mut self) {
 37        self.data.reset();
 38        self.basek = 0;
 39        self.k = 0;
 40        self.digits = 0;
 41        self.rank_out = 0;
 42        self.seed = Vec::new();
 43    }
 44}
 45
 46impl<D: DataRef> fmt::Display for GGLWECiphertextCompressed<D> {
 47    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 48        write!(
 49            f,
 50            "(GGLWECiphertextCompressed: basek={} k={} digits={}) {}",
 51            self.basek, self.k, self.digits, self.data
 52        )
 53    }
 54}
 55
 56impl GGLWECiphertextCompressed<Vec<u8>> {
 57    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize, rank_out: usize) -> Self {
 58        let size: usize = k.div_ceil(basek);
 59        debug_assert!(
 60            size > digits,
 61            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
 62            size,
 63            digits
 64        );
 65
 66        assert!(
 67            rows * digits <= size,
 68            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
 69            rows,
 70            digits,
 71            size
 72        );
 73
 74        Self {
 75            data: MatZnx::alloc(n, rows, rank_in, 1, size),
 76            basek: basek,
 77            k,
 78            rank_out,
 79            digits,
 80            seed: vec![[0u8; 32]; rows * rank_in],
 81        }
 82    }
 83
 84    pub fn bytes_of(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize) -> usize {
 85        let size: usize = k.div_ceil(basek);
 86        debug_assert!(
 87            size > digits,
 88            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
 89            size,
 90            digits
 91        );
 92
 93        assert!(
 94            rows * digits <= size,
 95            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
 96            rows,
 97            digits,
 98            size
 99        );
100
101        MatZnx::alloc_bytes(n, rows, rank_in, 1, rows)
102    }
103}
104
105impl<D: Data> Infos for GGLWECiphertextCompressed<D> {
106    type Inner = MatZnx<D>;
107
108    fn inner(&self) -> &Self::Inner {
109        &self.data
110    }
111
112    fn basek(&self) -> usize {
113        self.basek
114    }
115
116    fn k(&self) -> usize {
117        self.k
118    }
119}
120
121impl<D: Data> GGLWECiphertextCompressed<D> {
122    pub fn rank(&self) -> usize {
123        self.rank_out
124    }
125
126    pub fn digits(&self) -> usize {
127        self.digits
128    }
129
130    pub fn rank_in(&self) -> usize {
131        self.data.cols_in()
132    }
133
134    pub fn rank_out(&self) -> usize {
135        self.rank_out
136    }
137}
138
139impl<D: DataRef> GGLWECiphertextCompressed<D> {
140    pub(crate) fn at(&self, row: usize, col: usize) -> GLWECiphertextCompressed<&[u8]> {
141        GLWECiphertextCompressed {
142            data: self.data.at(row, col),
143            basek: self.basek,
144            k: self.k,
145            rank: self.rank_out,
146            seed: self.seed[self.rank_in() * row + col],
147        }
148    }
149}
150
151impl<D: DataMut> GGLWECiphertextCompressed<D> {
152    pub(crate) fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertextCompressed<&mut [u8]> {
153        let rank_in: usize = self.rank_in();
154        GLWECiphertextCompressed {
155            data: self.data.at_mut(row, col),
156            basek: self.basek,
157            k: self.k,
158            rank: self.rank_out,
159            seed: self.seed[rank_in * row + col], // Warning: value is copied and not borrow mut
160        }
161    }
162}
163
164impl<D: DataMut> ReaderFrom for GGLWECiphertextCompressed<D> {
165    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
166        self.k = reader.read_u64::<LittleEndian>()? as usize;
167        self.basek = reader.read_u64::<LittleEndian>()? as usize;
168        self.digits = reader.read_u64::<LittleEndian>()? as usize;
169        self.rank_out = reader.read_u64::<LittleEndian>()? as usize;
170        let seed_len = reader.read_u64::<LittleEndian>()? as usize;
171        self.seed = vec![[0u8; 32]; seed_len];
172        for s in &mut self.seed {
173            reader.read_exact(s)?;
174        }
175        self.data.read_from(reader)
176    }
177}
178
179impl<D: DataRef> WriterTo for GGLWECiphertextCompressed<D> {
180    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
181        writer.write_u64::<LittleEndian>(self.k as u64)?;
182        writer.write_u64::<LittleEndian>(self.basek as u64)?;
183        writer.write_u64::<LittleEndian>(self.digits as u64)?;
184        writer.write_u64::<LittleEndian>(self.rank_out as u64)?;
185        writer.write_u64::<LittleEndian>(self.seed.len() as u64)?;
186        for s in &self.seed {
187            writer.write_all(s)?;
188        }
189        self.data.write_to(writer)
190    }
191}
192
193impl<D: DataMut, B: Backend, DR: DataRef> Decompress<B, GGLWECiphertextCompressed<DR>> for GGLWECiphertext<D> {
194    fn decompress(&mut self, module: &Module<B>, other: &GGLWECiphertextCompressed<DR>)
195    where
196        Module<B>: VecZnxFillUniform + VecZnxCopy,
197    {
198        #[cfg(debug_assertions)]
199        {
200            use backend::hal::api::ZnxInfos;
201
202            assert_eq!(
203                self.n(),
204                other.data.n(),
205                "invalid receiver: self.n()={} != other.n()={}",
206                self.n(),
207                other.data.n()
208            );
209            assert_eq!(
210                self.size(),
211                other.size(),
212                "invalid receiver: self.size()={} != other.size()={}",
213                self.size(),
214                other.size()
215            );
216            assert_eq!(
217                self.rank_in(),
218                other.rank_in(),
219                "invalid receiver: self.rank_in()={} != other.rank_in()={}",
220                self.rank_in(),
221                other.rank_in()
222            );
223            assert_eq!(
224                self.rank_out(),
225                other.rank_out(),
226                "invalid receiver: self.rank_out()={} != other.rank_out()={}",
227                self.rank_out(),
228                other.rank_out()
229            );
230
231            assert_eq!(
232                self.rows(),
233                other.rows(),
234                "invalid receiver: self.rows()={} != other.rows()={}",
235                self.rows(),
236                other.rows()
237            );
238        }
239
240        let rank_in: usize = self.rank_in();
241        let rows: usize = self.rows();
242
243        (0..rank_in).for_each(|col_i| {
244            (0..rows).for_each(|row_i| {
245                self.at_mut(row_i, col_i)
246                    .decompress(module, &other.at(row_i, col_i));
247            });
248        });
249    }
250}
251
252#[derive(PartialEq, Eq, Clone)]
253pub struct GLWESwitchingKeyCompressed<D: Data> {
254    pub(crate) key: GGLWECiphertextCompressed<D>,
255    pub(crate) sk_in_n: usize,  // Degree of sk_in
256    pub(crate) sk_out_n: usize, // Degree of sk_out
257}
258
259impl<D: DataRef> fmt::Debug for GLWESwitchingKeyCompressed<D> {
260    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
261        write!(f, "{}", self)
262    }
263}
264
265impl<D: DataMut> FillUniform for GLWESwitchingKeyCompressed<D> {
266    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
267        self.key.fill_uniform(source);
268    }
269}
270
271impl<D: DataMut> Reset for GLWESwitchingKeyCompressed<D>
272where
273    MatZnx<D>: Reset,
274{
275    fn reset(&mut self) {
276        self.key.reset();
277        self.sk_in_n = 0;
278        self.sk_out_n = 0;
279    }
280}
281
282impl<D: DataRef> fmt::Display for GLWESwitchingKeyCompressed<D> {
283    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
284        write!(
285            f,
286            "(GLWESwitchingKeyCompressed: sk_in_n={} sk_out_n={}) {}",
287            self.sk_in_n, self.sk_out_n, self.key.data
288        )
289    }
290}
291
292impl<D: Data> Infos for GLWESwitchingKeyCompressed<D> {
293    type Inner = MatZnx<D>;
294
295    fn inner(&self) -> &Self::Inner {
296        &self.key.inner()
297    }
298
299    fn basek(&self) -> usize {
300        self.key.basek()
301    }
302
303    fn k(&self) -> usize {
304        self.key.k()
305    }
306}
307
308impl<D: Data> GLWESwitchingKeyCompressed<D> {
309    pub fn rank(&self) -> usize {
310        self.key.rank()
311    }
312
313    pub fn digits(&self) -> usize {
314        self.key.digits()
315    }
316
317    pub fn rank_in(&self) -> usize {
318        self.key.rank_in()
319    }
320
321    pub fn rank_out(&self) -> usize {
322        self.key.rank_out()
323    }
324}
325
326impl GLWESwitchingKeyCompressed<Vec<u8>> {
327    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize, rank_out: usize) -> Self {
328        GLWESwitchingKeyCompressed {
329            key: GGLWECiphertextCompressed::alloc(n, basek, k, rows, digits, rank_in, rank_out),
330            sk_in_n: 0,
331            sk_out_n: 0,
332        }
333    }
334
335    pub fn bytes_of(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize) -> usize {
336        GGLWECiphertextCompressed::bytes_of(n, basek, k, rows, digits, rank_in)
337    }
338}
339
340impl<D: DataMut> ReaderFrom for GLWESwitchingKeyCompressed<D> {
341    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
342        self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
343        self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
344        self.key.read_from(reader)
345    }
346}
347
348impl<D: DataRef> WriterTo for GLWESwitchingKeyCompressed<D> {
349    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
350        writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
351        writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
352        self.key.write_to(writer)
353    }
354}
355
356impl<D: DataMut> GLWESwitchingKey<D> {
357    pub fn decompress<DataOther: DataRef, B: Backend>(
358        &mut self,
359        module: &Module<B>,
360        other: &GLWESwitchingKeyCompressed<DataOther>,
361    ) where
362        Module<B>: VecZnxFillUniform + VecZnxCopy,
363    {
364        self.key.decompress(module, &other.key);
365        self.sk_in_n = other.sk_in_n;
366        self.sk_out_n = other.sk_out_n;
367    }
368}
369
370#[derive(PartialEq, Eq, Clone)]
371pub struct AutomorphismKeyCompressed<D: Data> {
372    pub(crate) key: GLWESwitchingKeyCompressed<D>,
373    pub(crate) p: i64,
374}
375
376impl<D: DataRef> fmt::Debug for AutomorphismKeyCompressed<D> {
377    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
378        write!(f, "{}", self)
379    }
380}
381
382impl<D: DataMut> FillUniform for AutomorphismKeyCompressed<D> {
383    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
384        self.key.fill_uniform(source);
385    }
386}
387
388impl<D: DataMut> Reset for AutomorphismKeyCompressed<D>
389where
390    MatZnx<D>: Reset,
391{
392    fn reset(&mut self) {
393        self.key.reset();
394        self.p = 0;
395    }
396}
397
398impl<D: DataRef> fmt::Display for AutomorphismKeyCompressed<D> {
399    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
400        write!(f, "(AutomorphismKeyCompressed: p={}) {}", self.p, self.key)
401    }
402}
403
404impl AutomorphismKeyCompressed<Vec<u8>> {
405    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self {
406        AutomorphismKeyCompressed {
407            key: GLWESwitchingKeyCompressed::alloc(n, basek, k, rows, digits, rank, rank),
408            p: 0,
409        }
410    }
411
412    pub fn bytes_of(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize {
413        GLWESwitchingKeyCompressed::<Vec<u8>>::bytes_of(n, basek, k, rows, digits, rank)
414    }
415}
416
417impl<D: Data> Infos for AutomorphismKeyCompressed<D> {
418    type Inner = MatZnx<D>;
419
420    fn inner(&self) -> &Self::Inner {
421        &self.key.inner()
422    }
423
424    fn basek(&self) -> usize {
425        self.key.basek()
426    }
427
428    fn k(&self) -> usize {
429        self.key.k()
430    }
431}
432
433impl<D: Data> AutomorphismKeyCompressed<D> {
434    pub fn rank(&self) -> usize {
435        self.key.rank()
436    }
437
438    pub fn digits(&self) -> usize {
439        self.key.digits()
440    }
441
442    pub fn rank_in(&self) -> usize {
443        self.key.rank_in()
444    }
445
446    pub fn rank_out(&self) -> usize {
447        self.key.rank_out()
448    }
449}
450
451impl<D: DataMut> ReaderFrom for AutomorphismKeyCompressed<D> {
452    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
453        self.p = reader.read_u64::<LittleEndian>()? as i64;
454        self.key.read_from(reader)
455    }
456}
457
458impl<D: DataRef> WriterTo for AutomorphismKeyCompressed<D> {
459    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
460        writer.write_u64::<LittleEndian>(self.p as u64)?;
461        self.key.write_to(writer)
462    }
463}
464
465impl<D: DataMut> AutomorphismKey<D> {
466    pub fn decompress<DataOther: DataRef, B: Backend>(&mut self, module: &Module<B>, other: &AutomorphismKeyCompressed<DataOther>)
467    where
468        Module<B>: VecZnxFillUniform + VecZnxCopy,
469    {
470        self.key.decompress(module, &other.key);
471        self.p = other.p;
472    }
473}
474
475#[derive(PartialEq, Eq, Clone)]
476pub struct GLWETensorKeyCompressed<D: Data> {
477    pub(crate) keys: Vec<GLWESwitchingKeyCompressed<D>>,
478}
479
480impl<D: DataRef> fmt::Debug for GLWETensorKeyCompressed<D> {
481    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
482        write!(f, "{}", self)
483    }
484}
485
486impl<D: DataMut> FillUniform for GLWETensorKeyCompressed<D> {
487    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
488        self.keys
489            .iter_mut()
490            .for_each(|key: &mut GLWESwitchingKeyCompressed<D>| key.fill_uniform(source))
491    }
492}
493
494impl<D: DataMut> Reset for GLWETensorKeyCompressed<D>
495where
496    MatZnx<D>: Reset,
497{
498    fn reset(&mut self) {
499        self.keys
500            .iter_mut()
501            .for_each(|key: &mut GLWESwitchingKeyCompressed<D>| key.reset())
502    }
503}
504
505impl<D: DataRef> fmt::Display for GLWETensorKeyCompressed<D> {
506    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
507        writeln!(f, "(GLWETensorKeyCompressed)",)?;
508        for (i, key) in self.keys.iter().enumerate() {
509            write!(f, "{}: {}", i, key)?;
510        }
511        Ok(())
512    }
513}
514
515impl GLWETensorKeyCompressed<Vec<u8>> {
516    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self {
517        let mut keys: Vec<GLWESwitchingKeyCompressed<Vec<u8>>> = Vec::new();
518        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
519        (0..pairs).for_each(|_| {
520            keys.push(GLWESwitchingKeyCompressed::alloc(
521                n, basek, k, rows, digits, 1, rank,
522            ));
523        });
524        Self { keys: keys }
525    }
526
527    pub fn bytes_of(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize {
528        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
529        pairs * GLWESwitchingKeyCompressed::bytes_of(n, basek, k, rows, digits, 1)
530    }
531}
532
533impl<D: Data> Infos for GLWETensorKeyCompressed<D> {
534    type Inner = MatZnx<D>;
535
536    fn inner(&self) -> &Self::Inner {
537        &self.keys[0].inner()
538    }
539
540    fn basek(&self) -> usize {
541        self.keys[0].basek()
542    }
543
544    fn k(&self) -> usize {
545        self.keys[0].k()
546    }
547}
548
549impl<D: Data> GLWETensorKeyCompressed<D> {
550    pub fn rank(&self) -> usize {
551        self.keys[0].rank()
552    }
553
554    pub fn digits(&self) -> usize {
555        self.keys[0].digits()
556    }
557
558    pub fn rank_in(&self) -> usize {
559        self.keys[0].rank_in()
560    }
561
562    pub fn rank_out(&self) -> usize {
563        self.keys[0].rank_out()
564    }
565}
566
567impl<D: DataMut> ReaderFrom for GLWETensorKeyCompressed<D> {
568    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
569        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
570        if self.keys.len() != len {
571            return Err(std::io::Error::new(
572                std::io::ErrorKind::InvalidData,
573                format!("self.keys.len()={} != read len={}", self.keys.len(), len),
574            ));
575        }
576        for key in &mut self.keys {
577            key.read_from(reader)?;
578        }
579        Ok(())
580    }
581}
582
583impl<D: DataRef> WriterTo for GLWETensorKeyCompressed<D> {
584    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
585        writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
586        for key in &self.keys {
587            key.write_to(writer)?;
588        }
589        Ok(())
590    }
591}
592
593impl<D: DataMut> GLWETensorKeyCompressed<D> {
594    pub(crate) fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GLWESwitchingKeyCompressed<D> {
595        if i > j {
596            std::mem::swap(&mut i, &mut j);
597        };
598        let rank: usize = self.rank();
599        &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
600    }
601}
602
603impl<D: DataMut> GLWETensorKey<D> {
604    pub fn decompress<DataOther: DataRef, B: Backend>(&mut self, module: &Module<B>, other: &GLWETensorKeyCompressed<DataOther>)
605    where
606        Module<B>: VecZnxFillUniform + VecZnxCopy,
607    {
608        #[cfg(debug_assertions)]
609        {
610            assert_eq!(
611                self.keys.len(),
612                other.keys.len(),
613                "invalid receiver: self.keys.len()={} != other.keys.len()={}",
614                self.keys.len(),
615                other.keys.len()
616            );
617        }
618
619        self.keys
620            .iter_mut()
621            .zip(other.keys.iter())
622            .for_each(|(a, b)| {
623                a.decompress(module, b);
624            });
625    }
626}