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