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