1use std::fmt;
  2
  3use backend::hal::{
  4    api::{FillUniform, Reset, VecZnxFillUniform, ZnxInfos, ZnxView, ZnxViewMut},
  5    layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, VecZnx, WriterTo},
  6};
  7use sampling::source::Source;
  8
  9use crate::{
 10    Decompress, GGLWEEncryptSkFamily, GLWESwitchingKeyCompressed, GLWEToLWESwitchingKey, Infos, LWECiphertext, LWESwitchingKey,
 11    LWEToGLWESwitchingKey, SetMetaData,
 12};
 13
 14#[derive(PartialEq, Eq, Clone)]
 15pub struct LWECiphertextCompressed<D: Data> {
 16    pub(crate) data: VecZnx<D>,
 17    pub(crate) k: usize,
 18    pub(crate) basek: usize,
 19    pub(crate) seed: [u8; 32],
 20}
 21
 22impl<D: DataRef> fmt::Debug for LWECiphertextCompressed<D> {
 23    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 24        write!(f, "{}", self)
 25    }
 26}
 27
 28impl<D: DataRef> fmt::Display for LWECiphertextCompressed<D> {
 29    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 30        write!(
 31            f,
 32            "LWECiphertextCompressed: basek={} k={} seed={:?}: {}",
 33            self.basek(),
 34            self.k(),
 35            self.seed,
 36            self.data
 37        )
 38    }
 39}
 40
 41impl<D: DataMut> Reset for LWECiphertextCompressed<D>
 42where
 43    VecZnx<D>: Reset,
 44{
 45    fn reset(&mut self) {
 46        self.data.reset();
 47        self.basek = 0;
 48        self.k = 0;
 49        self.seed = [0u8; 32];
 50    }
 51}
 52
 53impl<D: DataMut> FillUniform for LWECiphertextCompressed<D> {
 54    fn fill_uniform(&mut self, source: &mut Source) {
 55        self.data.fill_uniform(source);
 56    }
 57}
 58
 59impl LWECiphertextCompressed<Vec<u8>> {
 60    pub fn alloc(basek: usize, k: usize) -> Self {
 61        Self {
 62            data: VecZnx::alloc(1, 1, k.div_ceil(basek)),
 63            k: k,
 64            basek: basek,
 65            seed: [0u8; 32],
 66        }
 67    }
 68}
 69
 70impl<D: Data> Infos for LWECiphertextCompressed<D>
 71where
 72    VecZnx<D>: ZnxInfos,
 73{
 74    type Inner = VecZnx<D>;
 75
 76    fn n(&self) -> usize {
 77        &self.inner().n() - 1
 78    }
 79
 80    fn inner(&self) -> &Self::Inner {
 81        &self.data
 82    }
 83
 84    fn basek(&self) -> usize {
 85        self.basek
 86    }
 87
 88    fn k(&self) -> usize {
 89        self.k
 90    }
 91}
 92
 93impl<DataSelf: DataMut> SetMetaData for LWECiphertextCompressed<DataSelf> {
 94    fn set_k(&mut self, k: usize) {
 95        self.k = k
 96    }
 97
 98    fn set_basek(&mut self, basek: usize) {
 99        self.basek = basek
100    }
101}
102
103use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
104
105impl<D: DataMut> ReaderFrom for LWECiphertextCompressed<D> {
106    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
107        self.k = reader.read_u64::<LittleEndian>()? as usize;
108        self.basek = reader.read_u64::<LittleEndian>()? as usize;
109        reader.read(&mut self.seed)?;
110        self.data.read_from(reader)
111    }
112}
113
114impl<D: DataRef> WriterTo for LWECiphertextCompressed<D> {
115    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
116        writer.write_u64::<LittleEndian>(self.k as u64)?;
117        writer.write_u64::<LittleEndian>(self.basek as u64)?;
118        writer.write_all(&self.seed)?;
119        self.data.write_to(writer)
120    }
121}
122
123impl<D: DataMut, B: Backend, DR: DataRef> Decompress<B, LWECiphertextCompressed<DR>> for LWECiphertext<D> {
124    fn decompress(&mut self, module: &Module<B>, other: &LWECiphertextCompressed<DR>)
125    where
126        Module<B>: VecZnxFillUniform,
127    {
128        let mut source = Source::new(other.seed);
129        module.vec_znx_fill_uniform(other.basek(), &mut self.data, 0, other.k(), &mut source);
130        (0..self.size()).for_each(|i| {
131            self.data.at_mut(0, i)[0] = other.data.at(0, i)[0];
132        });
133    }
134}
135
136#[derive(PartialEq, Eq, Clone)]
137pub struct GLWEToLWESwitchingKeyCompressed<D: Data>(pub(crate) GLWESwitchingKeyCompressed<D>);
138
139impl<D: DataRef> fmt::Debug for GLWEToLWESwitchingKeyCompressed<D> {
140    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
141        write!(f, "{}", self)
142    }
143}
144
145impl<D: DataMut> FillUniform for GLWEToLWESwitchingKeyCompressed<D> {
146    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
147        self.0.fill_uniform(source);
148    }
149}
150
151impl<D: DataMut> Reset for GLWEToLWESwitchingKeyCompressed<D> {
152    fn reset(&mut self) {
153        self.0.reset();
154    }
155}
156
157impl<D: DataRef> fmt::Display for GLWEToLWESwitchingKeyCompressed<D> {
158    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
159        write!(f, "(GLWEToLWESwitchingKeyCompressed) {}", self.0)
160    }
161}
162
163impl<D: Data> Infos for GLWEToLWESwitchingKeyCompressed<D> {
164    type Inner = MatZnx<D>;
165
166    fn inner(&self) -> &Self::Inner {
167        &self.0.inner()
168    }
169
170    fn basek(&self) -> usize {
171        self.0.basek()
172    }
173
174    fn k(&self) -> usize {
175        self.0.k()
176    }
177}
178
179impl<D: Data> GLWEToLWESwitchingKeyCompressed<D> {
180    pub fn digits(&self) -> usize {
181        self.0.digits()
182    }
183
184    pub fn rank(&self) -> usize {
185        self.0.rank()
186    }
187
188    pub fn rank_in(&self) -> usize {
189        self.0.rank_in()
190    }
191
192    pub fn rank_out(&self) -> usize {
193        self.0.rank_out()
194    }
195}
196
197impl<D: DataMut> ReaderFrom for GLWEToLWESwitchingKeyCompressed<D> {
198    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
199        self.0.read_from(reader)
200    }
201}
202
203impl<D: DataRef> WriterTo for GLWEToLWESwitchingKeyCompressed<D> {
204    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
205        self.0.write_to(writer)
206    }
207}
208
209impl GLWEToLWESwitchingKeyCompressed<Vec<u8>> {
210    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, rank_in: usize) -> Self {
211        Self(GLWESwitchingKeyCompressed::alloc(
212            n, basek, k, rows, 1, rank_in, 1,
213        ))
214    }
215
216    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank_in: usize) -> usize
217    where
218        Module<B>: GGLWEEncryptSkFamily<B>,
219    {
220        GLWEToLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, rank_in)
221    }
222}
223
224#[derive(PartialEq, Eq, Clone)]
225pub struct LWEToGLWESwitchingKeyCompressed<D: Data>(pub(crate) GLWESwitchingKeyCompressed<D>);
226
227impl<D: DataRef> fmt::Debug for LWEToGLWESwitchingKeyCompressed<D> {
228    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
229        write!(f, "{}", self)
230    }
231}
232
233impl<D: DataMut> FillUniform for LWEToGLWESwitchingKeyCompressed<D> {
234    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
235        self.0.fill_uniform(source);
236    }
237}
238
239impl<D: DataMut> Reset for LWEToGLWESwitchingKeyCompressed<D> {
240    fn reset(&mut self) {
241        self.0.reset();
242    }
243}
244
245impl<D: DataRef> fmt::Display for LWEToGLWESwitchingKeyCompressed<D> {
246    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
247        write!(f, "(LWEToGLWESwitchingKeyCompressed) {}", self.0)
248    }
249}
250
251impl<D: Data> Infos for LWEToGLWESwitchingKeyCompressed<D> {
252    type Inner = MatZnx<D>;
253
254    fn inner(&self) -> &Self::Inner {
255        &self.0.inner()
256    }
257
258    fn basek(&self) -> usize {
259        self.0.basek()
260    }
261
262    fn k(&self) -> usize {
263        self.0.k()
264    }
265}
266
267impl<D: Data> LWEToGLWESwitchingKeyCompressed<D> {
268    pub fn digits(&self) -> usize {
269        self.0.digits()
270    }
271
272    pub fn rank(&self) -> usize {
273        self.0.rank()
274    }
275
276    pub fn rank_in(&self) -> usize {
277        self.0.rank_in()
278    }
279
280    pub fn rank_out(&self) -> usize {
281        self.0.rank_out()
282    }
283}
284
285impl<D: DataMut> ReaderFrom for LWEToGLWESwitchingKeyCompressed<D> {
286    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
287        self.0.read_from(reader)
288    }
289}
290
291impl<D: DataRef> WriterTo for LWEToGLWESwitchingKeyCompressed<D> {
292    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
293        self.0.write_to(writer)
294    }
295}
296
297impl LWEToGLWESwitchingKeyCompressed<Vec<u8>> {
298    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, rank_out: usize) -> Self {
299        Self(GLWESwitchingKeyCompressed::alloc(
300            n, basek, k, rows, 1, 1, rank_out,
301        ))
302    }
303
304    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize, rank_out: usize) -> usize
305    where
306        Module<B>: GGLWEEncryptSkFamily<B>,
307    {
308        LWEToGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k, rank_out)
309    }
310}
311
312#[derive(PartialEq, Eq, Clone)]
313pub struct LWESwitchingKeyCompressed<D: Data>(pub(crate) GLWESwitchingKeyCompressed<D>);
314
315impl<D: DataRef> fmt::Debug for LWESwitchingKeyCompressed<D> {
316    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
317        write!(f, "{}", self)
318    }
319}
320
321impl<D: DataMut> FillUniform for LWESwitchingKeyCompressed<D> {
322    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
323        self.0.fill_uniform(source);
324    }
325}
326
327impl<D: DataMut> Reset for LWESwitchingKeyCompressed<D> {
328    fn reset(&mut self) {
329        self.0.reset();
330    }
331}
332
333impl<D: DataRef> fmt::Display for LWESwitchingKeyCompressed<D> {
334    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
335        write!(f, "(LWESwitchingKeyCompressed) {}", self.0)
336    }
337}
338
339impl<D: Data> Infos for LWESwitchingKeyCompressed<D> {
340    type Inner = MatZnx<D>;
341
342    fn inner(&self) -> &Self::Inner {
343        &self.0.inner()
344    }
345
346    fn basek(&self) -> usize {
347        self.0.basek()
348    }
349
350    fn k(&self) -> usize {
351        self.0.k()
352    }
353}
354
355impl<D: Data> LWESwitchingKeyCompressed<D> {
356    pub fn digits(&self) -> usize {
357        self.0.digits()
358    }
359
360    pub fn rank(&self) -> usize {
361        self.0.rank()
362    }
363
364    pub fn rank_in(&self) -> usize {
365        self.0.rank_in()
366    }
367
368    pub fn rank_out(&self) -> usize {
369        self.0.rank_out()
370    }
371}
372
373impl<D: DataMut> ReaderFrom for LWESwitchingKeyCompressed<D> {
374    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
375        self.0.read_from(reader)
376    }
377}
378
379impl<D: DataRef> WriterTo for LWESwitchingKeyCompressed<D> {
380    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
381        self.0.write_to(writer)
382    }
383}
384
385impl LWESwitchingKeyCompressed<Vec<u8>> {
386    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize) -> Self {
387        Self(GLWESwitchingKeyCompressed::alloc(
388            n, basek, k, rows, 1, 1, 1,
389        ))
390    }
391
392    pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
393    where
394        Module<B>: GGLWEEncryptSkFamily<B>,
395    {
396        LWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k)
397    }
398}