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