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::{Base2K, Degree, GLWECiphertext, GLWEInfos, LWEInfos, Rank, TorusPrecision, compressed::Decompress};
 10use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 11use std::fmt;
 12
 13#[derive(PartialEq, Eq, Clone)]
 14pub struct GLWECiphertextCompressed<D: Data> {
 15    pub(crate) data: VecZnx<D>,
 16    pub(crate) base2k: Base2K,
 17    pub(crate) k: TorusPrecision,
 18    pub(crate) rank: Rank,
 19    pub(crate) seed: [u8; 32],
 20}
 21
 22impl<D: Data> LWEInfos for GLWECiphertextCompressed<D> {
 23    fn base2k(&self) -> Base2K {
 24        self.base2k
 25    }
 26
 27    fn k(&self) -> TorusPrecision {
 28        self.k
 29    }
 30
 31    fn size(&self) -> usize {
 32        self.data.size()
 33    }
 34
 35    fn n(&self) -> Degree {
 36        Degree(self.data.n() as u32)
 37    }
 38}
 39impl<D: Data> GLWEInfos for GLWECiphertextCompressed<D> {
 40    fn rank(&self) -> Rank {
 41        self.rank
 42    }
 43}
 44
 45impl<D: DataRef> fmt::Debug for GLWECiphertextCompressed<D> {
 46    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 47        write!(f, "{self}")
 48    }
 49}
 50
 51impl<D: DataRef> fmt::Display for GLWECiphertextCompressed<D> {
 52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 53        write!(
 54            f,
 55            "GLWECiphertextCompressed: base2k={} k={} rank={} seed={:?}: {}",
 56            self.base2k(),
 57            self.k(),
 58            self.rank(),
 59            self.seed,
 60            self.data
 61        )
 62    }
 63}
 64
 65impl<D: DataMut> FillUniform for GLWECiphertextCompressed<D> {
 66    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
 67        self.data.fill_uniform(log_bound, source);
 68    }
 69}
 70
 71impl GLWECiphertextCompressed<Vec<u8>> {
 72    pub fn alloc<A>(infos: &A) -> Self
 73    where
 74        A: GLWEInfos,
 75    {
 76        Self::alloc_with(infos.n(), infos.base2k(), infos.k(), infos.rank())
 77    }
 78
 79    pub fn alloc_with(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self {
 80        Self {
 81            data: VecZnx::alloc(n.into(), 1, k.0.div_ceil(base2k.0) as usize),
 82            base2k,
 83            k,
 84            rank,
 85            seed: [0u8; 32],
 86        }
 87    }
 88
 89    pub fn alloc_bytes<A>(infos: &A) -> usize
 90    where
 91        A: GLWEInfos,
 92    {
 93        Self::alloc_bytes_with(infos.n(), infos.base2k(), infos.k())
 94    }
 95
 96    pub fn alloc_bytes_with(n: Degree, base2k: Base2K, k: TorusPrecision) -> usize {
 97        VecZnx::alloc_bytes(n.into(), 1, k.0.div_ceil(base2k.0) as usize)
 98    }
 99}
100
101impl<D: DataMut> ReaderFrom for GLWECiphertextCompressed<D> {
102    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
103        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
104        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
105        self.rank = Rank(reader.read_u32::<LittleEndian>()?);
106        reader.read_exact(&mut self.seed)?;
107        self.data.read_from(reader)
108    }
109}
110
111impl<D: DataRef> WriterTo for GLWECiphertextCompressed<D> {
112    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
113        writer.write_u32::<LittleEndian>(self.k.into())?;
114        writer.write_u32::<LittleEndian>(self.base2k.into())?;
115        writer.write_u32::<LittleEndian>(self.rank.into())?;
116        writer.write_all(&self.seed)?;
117        self.data.write_to(writer)
118    }
119}
120
121impl<D: DataMut, B: Backend, DR: DataRef> Decompress<B, GLWECiphertextCompressed<DR>> for GLWECiphertext<D>
122where
123    Module<B>: VecZnxFillUniform + VecZnxCopy,
124{
125    fn decompress(&mut self, module: &Module<B>, other: &GLWECiphertextCompressed<DR>) {
126        #[cfg(debug_assertions)]
127        {
128            assert_eq!(
129                self.n(),
130                other.n(),
131                "invalid receiver: self.n()={} != other.n()={}",
132                self.n(),
133                other.n()
134            );
135            assert_eq!(
136                self.size(),
137                other.size(),
138                "invalid receiver: self.size()={} != other.size()={}",
139                self.size(),
140                other.size()
141            );
142            assert_eq!(
143                self.rank(),
144                other.rank(),
145                "invalid receiver: self.rank()={} != other.rank()={}",
146                self.rank(),
147                other.rank()
148            );
149        }
150
151        let mut source: Source = Source::new(other.seed);
152        self.decompress_internal(module, other, &mut source);
153    }
154}
155
156impl<D: DataMut> GLWECiphertext<D> {
157    pub(crate) fn decompress_internal<DataOther, B: Backend>(
158        &mut self,
159        module: &Module<B>,
160        other: &GLWECiphertextCompressed<DataOther>,
161        source: &mut Source,
162    ) where
163        DataOther: DataRef,
164        Module<B>: VecZnxCopy + VecZnxFillUniform,
165    {
166        #[cfg(debug_assertions)]
167        {
168            assert_eq!(self.rank(), other.rank());
169            debug_assert_eq!(self.size(), other.size());
170        }
171
172        module.vec_znx_copy(&mut self.data, 0, &other.data, 0);
173        (1..(other.rank() + 1).into()).for_each(|i| {
174            module.vec_znx_fill_uniform(other.base2k.into(), &mut self.data, i, source);
175        });
176
177        self.base2k = other.base2k;
178        self.k = other.k;
179    }
180}
181
182pub trait GLWECiphertextCompressedToRef {
183    fn to_ref(&self) -> GLWECiphertextCompressed<&[u8]>;
184}
185
186impl<D: DataRef> GLWECiphertextCompressedToRef for GLWECiphertextCompressed<D> {
187    fn to_ref(&self) -> GLWECiphertextCompressed<&[u8]> {
188        GLWECiphertextCompressed {
189            seed: self.seed.clone(),
190            base2k: self.base2k,
191            k: self.k,
192            rank: self.rank,
193            data: self.data.to_ref(),
194        }
195    }
196}
197
198pub trait GLWECiphertextCompressedToMut {
199    fn to_mut(&mut self) -> GLWECiphertextCompressed<&mut [u8]>;
200}
201
202impl<D: DataMut> GLWECiphertextCompressedToMut for GLWECiphertextCompressed<D> {
203    fn to_mut(&mut self) -> GLWECiphertextCompressed<&mut [u8]> {
204        GLWECiphertextCompressed {
205            seed: self.seed.clone(),
206            base2k: self.base2k,
207            k: self.k,
208            rank: self.rank,
209            data: self.data.to_mut(),
210        }
211    }
212}