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