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