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, GLWE, GLWEInfos, GLWEToMut, GetDegree, LWEInfos, Rank, SetGLWEInfos, TorusPrecision};
 10use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 11use std::fmt;
 12
 13#[derive(PartialEq, Eq, Clone)]
 14pub struct GLWECompressed<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 GLWECompressed<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 GLWECompressed<D> {
 40    fn rank(&self) -> Rank {
 41        self.rank
 42    }
 43}
 44
 45impl<D: DataRef> fmt::Debug for GLWECompressed<D> {
 46    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 47        write!(f, "{self}")
 48    }
 49}
 50
 51impl<D: DataRef> fmt::Display for GLWECompressed<D> {
 52    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 53        write!(
 54            f,
 55            "GLWECompressed: 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 GLWECompressed<D> {
 66    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
 67        self.data.fill_uniform(log_bound, source);
 68    }
 69}
 70
 71pub trait GLWECompressedAlloc
 72where
 73    Self: GetDegree,
 74{
 75    fn alloc_glwe_compressed(&self, base2k: Base2K, k: TorusPrecision, rank: Rank) -> GLWECompressed<Vec<u8>> {
 76        GLWECompressed {
 77            data: VecZnx::alloc(self.n().into(), 1, k.0.div_ceil(base2k.0) as usize),
 78            base2k,
 79            k,
 80            rank,
 81            seed: [0u8; 32],
 82        }
 83    }
 84
 85    fn alloc_glwe_compressed_from_infos<A>(&self, infos: &A) -> GLWECompressed<Vec<u8>>
 86    where
 87        A: GLWEInfos,
 88    {
 89        assert_eq!(self.n(), infos.n());
 90        self.alloc_glwe_compressed(infos.base2k(), infos.k(), infos.rank())
 91    }
 92
 93    fn bytes_of_glwe_compressed(&self, base2k: Base2K, k: TorusPrecision) -> usize {
 94        VecZnx::bytes_of(self.n().into(), 1, k.0.div_ceil(base2k.0) as usize)
 95    }
 96
 97    fn bytes_of_glwe_compressed_from_infos<A>(&self, infos: &A) -> usize
 98    where
 99        A: GLWEInfos,
100    {
101        assert_eq!(self.n(), infos.n());
102        self.bytes_of_glwe_compressed(infos.base2k(), infos.k())
103    }
104}
105
106impl<B: Backend> GLWECompressedAlloc for Module<B> where Self: GetDegree {}
107
108impl GLWECompressed<Vec<u8>> {
109    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
110    where
111        A: GLWEInfos,
112        M: GLWECompressedAlloc,
113    {
114        module.alloc_glwe_compressed_from_infos(infos)
115    }
116
117    pub fn alloc<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank) -> Self
118    where
119        M: GLWECompressedAlloc,
120    {
121        module.alloc_glwe_compressed(base2k, k, rank)
122    }
123
124    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
125    where
126        A: GLWEInfos,
127        M: GLWECompressedAlloc,
128    {
129        module.bytes_of_glwe_compressed_from_infos(infos)
130    }
131
132    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision) -> usize
133    where
134        M: GLWECompressedAlloc,
135    {
136        module.bytes_of_glwe_compressed(base2k, k)
137    }
138}
139
140impl<D: DataMut> ReaderFrom for GLWECompressed<D> {
141    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
142        self.k = TorusPrecision(reader.read_u32::<LittleEndian>()?);
143        self.base2k = Base2K(reader.read_u32::<LittleEndian>()?);
144        self.rank = Rank(reader.read_u32::<LittleEndian>()?);
145        reader.read_exact(&mut self.seed)?;
146        self.data.read_from(reader)
147    }
148}
149
150impl<D: DataRef> WriterTo for GLWECompressed<D> {
151    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
152        writer.write_u32::<LittleEndian>(self.k.into())?;
153        writer.write_u32::<LittleEndian>(self.base2k.into())?;
154        writer.write_u32::<LittleEndian>(self.rank.into())?;
155        writer.write_all(&self.seed)?;
156        self.data.write_to(writer)
157    }
158}
159
160pub trait GLWEDecompress
161where
162    Self: GetDegree + VecZnxFillUniform + VecZnxCopy,
163{
164    fn decompress_glwe<R, O>(&self, res: &mut R, other: &O)
165    where
166        R: GLWEToMut + SetGLWEInfos,
167        O: GLWECompressedToRef + GLWEInfos,
168    {
169        {
170            let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
171            let other: &GLWECompressed<&[u8]> = &other.to_ref();
172            assert_eq!(
173                res.n(),
174                self.n(),
175                "invalid receiver: res.n()={} != other.n()={}",
176                res.n(),
177                self.n()
178            );
179
180            assert_eq!(res.lwe_layout(), other.lwe_layout());
181            assert_eq!(res.glwe_layout(), other.glwe_layout());
182
183            let mut source: Source = Source::new(other.seed);
184
185            self.vec_znx_copy(&mut res.data, 0, &other.data, 0);
186            (1..(other.rank() + 1).into()).for_each(|i| {
187                self.vec_znx_fill_uniform(other.base2k.into(), &mut res.data, i, &mut source);
188            });
189        }
190
191        res.set_base2k(other.base2k());
192        res.set_k(other.k());
193    }
194}
195
196impl<B: Backend> GLWEDecompress for Module<B> where Self: GetDegree + VecZnxFillUniform + VecZnxCopy {}
197
198impl<D: DataMut> GLWE<D> {
199    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
200    where
201        O: GLWECompressedToRef + GLWEInfos,
202        M: GLWEDecompress,
203    {
204        module.decompress_glwe(self, other);
205    }
206}
207
208pub trait GLWECompressedToRef {
209    fn to_ref(&self) -> GLWECompressed<&[u8]>;
210}
211
212impl<D: DataRef> GLWECompressedToRef for GLWECompressed<D> {
213    fn to_ref(&self) -> GLWECompressed<&[u8]> {
214        GLWECompressed {
215            seed: self.seed.clone(),
216            base2k: self.base2k,
217            k: self.k,
218            rank: self.rank,
219            data: self.data.to_ref(),
220        }
221    }
222}
223
224pub trait GLWECompressedToMut {
225    fn to_mut(&mut self) -> GLWECompressed<&mut [u8]>;
226}
227
228impl<D: DataMut> GLWECompressedToMut for GLWECompressed<D> {
229    fn to_mut(&mut self) -> GLWECompressed<&mut [u8]> {
230        GLWECompressed {
231            seed: self.seed.clone(),
232            base2k: self.base2k,
233            k: self.k,
234            rank: self.rank,
235            data: self.data.to_mut(),
236        }
237    }
238}