1use backend::hal::{
  2    api::{FillUniform, Reset, VecZnxCopy, VecZnxFillUniform, ZnxInfos},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, ReaderFrom, VecZnx, WriterTo},
  4};
  5use sampling::source::Source;
  6
  7use crate::{
  8    layouts::{GLWECiphertext, Infos},
  9    trait_families::{Decompress, DecompressFamily},
 10};
 11use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 12use std::fmt;
 13
 14#[derive(PartialEq, Eq, Clone)]
 15pub struct GLWECiphertextCompressed<D: Data> {
 16    pub(crate) data: VecZnx<D>,
 17    pub(crate) basek: usize,
 18    pub(crate) k: usize,
 19    pub(crate) rank: usize,
 20    pub(crate) seed: [u8; 32],
 21}
 22
 23impl<D: DataRef> fmt::Debug for GLWECiphertextCompressed<D> {
 24    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 25        write!(f, "{}", self)
 26    }
 27}
 28
 29impl<D: DataRef> fmt::Display for GLWECiphertextCompressed<D> {
 30    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 31        write!(
 32            f,
 33            "GLWECiphertextCompressed: basek={} k={} rank={} seed={:?}: {}",
 34            self.basek(),
 35            self.k(),
 36            self.rank,
 37            self.seed,
 38            self.data
 39        )
 40    }
 41}
 42
 43impl<D: DataMut> Reset for GLWECiphertextCompressed<D> {
 44    fn reset(&mut self) {
 45        self.data.reset();
 46        self.basek = 0;
 47        self.k = 0;
 48        self.rank = 0;
 49        self.seed = [0u8; 32];
 50    }
 51}
 52
 53impl<D: DataMut> FillUniform for GLWECiphertextCompressed<D> {
 54    fn fill_uniform(&mut self, source: &mut Source) {
 55        self.data.fill_uniform(source);
 56    }
 57}
 58
 59impl<D: Data> Infos for GLWECiphertextCompressed<D> {
 60    type Inner = VecZnx<D>;
 61
 62    fn inner(&self) -> &Self::Inner {
 63        &self.data
 64    }
 65
 66    fn basek(&self) -> usize {
 67        self.basek
 68    }
 69
 70    fn k(&self) -> usize {
 71        self.k
 72    }
 73}
 74
 75impl<D: Data> GLWECiphertextCompressed<D> {
 76    pub fn rank(&self) -> usize {
 77        self.rank
 78    }
 79}
 80
 81impl GLWECiphertextCompressed<Vec<u8>> {
 82    pub fn alloc(n: usize, basek: usize, k: usize, rank: usize) -> Self {
 83        Self {
 84            data: VecZnx::alloc(n, 1, k.div_ceil(basek)),
 85            basek,
 86            k,
 87            rank,
 88            seed: [0u8; 32],
 89        }
 90    }
 91
 92    pub fn bytes_of(n: usize, basek: usize, k: usize) -> usize {
 93        GLWECiphertext::bytes_of(n, basek, k, 1)
 94    }
 95}
 96
 97impl<D: DataMut> ReaderFrom for GLWECiphertextCompressed<D> {
 98    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
 99        self.k = reader.read_u64::<LittleEndian>()? as usize;
100        self.basek = reader.read_u64::<LittleEndian>()? as usize;
101        self.rank = reader.read_u64::<LittleEndian>()? as usize;
102        reader.read(&mut self.seed)?;
103        self.data.read_from(reader)
104    }
105}
106
107impl<D: DataRef> WriterTo for GLWECiphertextCompressed<D> {
108    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
109        writer.write_u64::<LittleEndian>(self.k as u64)?;
110        writer.write_u64::<LittleEndian>(self.basek as u64)?;
111        writer.write_u64::<LittleEndian>(self.rank as u64)?;
112        writer.write_all(&self.seed)?;
113        self.data.write_to(writer)
114    }
115}
116
117impl<D: DataMut, B: Backend, DR: DataRef> Decompress<B, GLWECiphertextCompressed<DR>> for GLWECiphertext<D> {
118    fn decompress(&mut self, module: &Module<B>, other: &GLWECiphertextCompressed<DR>)
119    where
120        Module<B>: DecompressFamily<B>,
121    {
122        #[cfg(debug_assertions)]
123        {
124            assert_eq!(
125                self.n(),
126                other.data.n(),
127                "invalid receiver: self.n()={} != other.n()={}",
128                self.n(),
129                other.data.n()
130            );
131            assert_eq!(
132                self.size(),
133                other.size(),
134                "invalid receiver: self.size()={} != other.size()={}",
135                self.size(),
136                other.size()
137            );
138            assert_eq!(
139                self.rank(),
140                other.rank(),
141                "invalid receiver: self.rank()={} != other.rank()={}",
142                self.rank(),
143                other.rank()
144            );
145            let mut source: Source = Source::new(other.seed);
146            self.decompress_internal(module, other, &mut source);
147        }
148    }
149}
150
151impl<D: DataMut> GLWECiphertext<D> {
152    pub(crate) fn decompress_internal<DataOther, B: Backend>(
153        &mut self,
154        module: &Module<B>,
155        other: &GLWECiphertextCompressed<DataOther>,
156        source: &mut Source,
157    ) where
158        DataOther: DataRef,
159        Module<B>: DecompressFamily<B>,
160    {
161        #[cfg(debug_assertions)]
162        {
163            assert_eq!(self.rank(), other.rank())
164        }
165
166        let k: usize = other.k;
167        let basek: usize = other.basek;
168        let cols: usize = other.rank() + 1;
169        module.vec_znx_copy(&mut self.data, 0, &other.data, 0);
170        (1..cols).for_each(|i| {
171            module.vec_znx_fill_uniform(basek, &mut self.data, i, k, source);
172        });
173
174        self.basek = basek;
175        self.k = k;
176    }
177}