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