1use std::fmt;
  2
  3use poulpy_hal::{
  4    api::ZnFillUniform,
  5    layouts::{
  6        Backend, Data, DataMut, DataRef, FillUniform, Module, ReaderFrom, Reset, VecZnx, WriterTo, ZnxInfos, ZnxView, ZnxViewMut,
  7    },
  8    source::Source,
  9};
 10
 11use crate::layouts::{Infos, LWECiphertext, SetMetaData, compressed::Decompress};
 12
 13#[derive(PartialEq, Eq, Clone)]
 14pub struct LWECiphertextCompressed<D: Data> {
 15    pub(crate) data: VecZnx<D>,
 16    pub(crate) k: usize,
 17    pub(crate) basek: usize,
 18    pub(crate) seed: [u8; 32],
 19}
 20
 21impl<D: DataRef> fmt::Debug for LWECiphertextCompressed<D> {
 22    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 23        write!(f, "{}", self)
 24    }
 25}
 26
 27impl<D: DataRef> fmt::Display for LWECiphertextCompressed<D> {
 28    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 29        write!(
 30            f,
 31            "LWECiphertextCompressed: basek={} k={} seed={:?}: {}",
 32            self.basek(),
 33            self.k(),
 34            self.seed,
 35            self.data
 36        )
 37    }
 38}
 39
 40impl<D: DataMut> Reset for LWECiphertextCompressed<D>
 41where
 42    VecZnx<D>: Reset,
 43{
 44    fn reset(&mut self) {
 45        self.data.reset();
 46        self.basek = 0;
 47        self.k = 0;
 48        self.seed = [0u8; 32];
 49    }
 50}
 51
 52impl<D: DataMut> FillUniform for LWECiphertextCompressed<D> {
 53    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
 54        self.data.fill_uniform(log_bound, source);
 55    }
 56}
 57
 58impl LWECiphertextCompressed<Vec<u8>> {
 59    pub fn alloc(basek: usize, k: usize) -> Self {
 60        Self {
 61            data: VecZnx::alloc(1, 1, k.div_ceil(basek)),
 62            k,
 63            basek,
 64            seed: [0u8; 32],
 65        }
 66    }
 67}
 68
 69impl<D: Data> Infos for LWECiphertextCompressed<D>
 70where
 71    VecZnx<D>: ZnxInfos,
 72{
 73    type Inner = VecZnx<D>;
 74
 75    fn n(&self) -> usize {
 76        &self.inner().n() - 1
 77    }
 78
 79    fn inner(&self) -> &Self::Inner {
 80        &self.data
 81    }
 82
 83    fn basek(&self) -> usize {
 84        self.basek
 85    }
 86
 87    fn k(&self) -> usize {
 88        self.k
 89    }
 90}
 91
 92impl<DataSelf: DataMut> SetMetaData for LWECiphertextCompressed<DataSelf> {
 93    fn set_k(&mut self, k: usize) {
 94        self.k = k
 95    }
 96
 97    fn set_basek(&mut self, basek: usize) {
 98        self.basek = basek
 99    }
100}
101
102use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
103
104impl<D: DataMut> ReaderFrom for LWECiphertextCompressed<D> {
105    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
106        self.k = reader.read_u64::<LittleEndian>()? as usize;
107        self.basek = reader.read_u64::<LittleEndian>()? as usize;
108        reader.read_exact(&mut self.seed)?;
109        self.data.read_from(reader)
110    }
111}
112
113impl<D: DataRef> WriterTo for LWECiphertextCompressed<D> {
114    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
115        writer.write_u64::<LittleEndian>(self.k as u64)?;
116        writer.write_u64::<LittleEndian>(self.basek as u64)?;
117        writer.write_all(&self.seed)?;
118        self.data.write_to(writer)
119    }
120}
121
122impl<D: DataMut, B: Backend, DR: DataRef> Decompress<B, LWECiphertextCompressed<DR>> for LWECiphertext<D>
123where
124    Module<B>: ZnFillUniform,
125{
126    fn decompress(&mut self, module: &Module<B>, other: &LWECiphertextCompressed<DR>) {
127        debug_assert_eq!(self.size(), other.size());
128        let mut source: Source = Source::new(other.seed);
129        module.zn_fill_uniform(self.n(), other.basek(), &mut self.data, 0, &mut source);
130        (0..self.size()).for_each(|i| {
131            self.data.at_mut(0, i)[0] = other.data.at(0, i)[0];
132        });
133    }
134}