1use backend::hal::{
  2    api::{FillUniform, Reset, VmpPMatAlloc, VmpPMatAllocBytes, VmpPMatPrepare},
  3    layouts::{Backend, Data, DataMut, DataRef, MatZnx, ReaderFrom, WriterTo},
  4};
  5
  6use crate::{GLWECiphertext, Infos};
  7use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
  8
  9pub trait GGLWEExecLayoutFamily<B: Backend> = VmpPMatAlloc<B> + VmpPMatAllocBytes + VmpPMatPrepare<B>;
 10use std::fmt;
 11
 12#[derive(PartialEq, Eq, Clone)]
 13pub struct GGLWECiphertext<D: Data> {
 14    pub(crate) data: MatZnx<D>,
 15    pub(crate) basek: usize,
 16    pub(crate) k: usize,
 17    pub(crate) digits: usize,
 18}
 19
 20impl<D: DataRef> fmt::Debug for GGLWECiphertext<D> {
 21    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 22        write!(f, "{}", self)
 23    }
 24}
 25
 26impl<D: DataMut> FillUniform for GGLWECiphertext<D> {
 27    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
 28        self.data.fill_uniform(source);
 29    }
 30}
 31
 32impl<D: DataMut> Reset for GGLWECiphertext<D>
 33where
 34    MatZnx<D>: Reset,
 35{
 36    fn reset(&mut self) {
 37        self.data.reset();
 38        self.basek = 0;
 39        self.k = 0;
 40        self.digits = 0;
 41    }
 42}
 43
 44impl<D: DataRef> fmt::Display for GGLWECiphertext<D> {
 45    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 46        write!(
 47            f,
 48            "(GGLWECiphertext: basek={} k={} digits={}) {}",
 49            self.basek, self.k, self.digits, self.data
 50        )
 51    }
 52}
 53
 54impl<D: DataRef> GGLWECiphertext<D> {
 55    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
 56        GLWECiphertext {
 57            data: self.data.at(row, col),
 58            basek: self.basek,
 59            k: self.k,
 60        }
 61    }
 62}
 63
 64impl<D: DataMut> GGLWECiphertext<D> {
 65    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
 66        GLWECiphertext {
 67            data: self.data.at_mut(row, col),
 68            basek: self.basek,
 69            k: self.k,
 70        }
 71    }
 72}
 73
 74impl GGLWECiphertext<Vec<u8>> {
 75    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize, rank_out: usize) -> Self {
 76        let size: usize = k.div_ceil(basek);
 77        debug_assert!(
 78            size > digits,
 79            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
 80            size,
 81            digits
 82        );
 83
 84        assert!(
 85            rows * digits <= size,
 86            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
 87            rows,
 88            digits,
 89            size
 90        );
 91
 92        Self {
 93            data: MatZnx::alloc(n, rows, rank_in, rank_out + 1, size),
 94            basek: basek,
 95            k,
 96            digits,
 97        }
 98    }
 99
100    pub fn bytes_of(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize, rank_out: usize) -> usize {
101        let size: usize = k.div_ceil(basek);
102        debug_assert!(
103            size > digits,
104            "invalid gglwe: ceil(k/basek): {} <= digits: {}",
105            size,
106            digits
107        );
108
109        assert!(
110            rows * digits <= size,
111            "invalid gglwe: rows: {} * digits:{} > ceil(k/basek): {}",
112            rows,
113            digits,
114            size
115        );
116
117        MatZnx::alloc_bytes(n, rows, rank_in, rank_out + 1, rows)
118    }
119}
120
121impl<D: Data> Infos for GGLWECiphertext<D> {
122    type Inner = MatZnx<D>;
123
124    fn inner(&self) -> &Self::Inner {
125        &self.data
126    }
127
128    fn basek(&self) -> usize {
129        self.basek
130    }
131
132    fn k(&self) -> usize {
133        self.k
134    }
135}
136
137impl<D: Data> GGLWECiphertext<D> {
138    pub fn rank(&self) -> usize {
139        self.data.cols_out() - 1
140    }
141
142    pub fn digits(&self) -> usize {
143        self.digits
144    }
145
146    pub fn rank_in(&self) -> usize {
147        self.data.cols_in()
148    }
149
150    pub fn rank_out(&self) -> usize {
151        self.data.cols_out() - 1
152    }
153}
154
155impl<D: DataMut> ReaderFrom for GGLWECiphertext<D> {
156    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
157        self.k = reader.read_u64::<LittleEndian>()? as usize;
158        self.basek = reader.read_u64::<LittleEndian>()? as usize;
159        self.digits = reader.read_u64::<LittleEndian>()? as usize;
160        self.data.read_from(reader)
161    }
162}
163
164impl<D: DataRef> WriterTo for GGLWECiphertext<D> {
165    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
166        writer.write_u64::<LittleEndian>(self.k as u64)?;
167        writer.write_u64::<LittleEndian>(self.basek as u64)?;
168        writer.write_u64::<LittleEndian>(self.digits as u64)?;
169        self.data.write_to(writer)
170    }
171}
172
173#[derive(PartialEq, Eq, Clone)]
174pub struct GLWESwitchingKey<D: Data> {
175    pub(crate) key: GGLWECiphertext<D>,
176    pub(crate) sk_in_n: usize,  // Degree of sk_in
177    pub(crate) sk_out_n: usize, // Degree of sk_out
178}
179
180impl<D: DataRef> fmt::Debug for GLWESwitchingKey<D> {
181    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
182        write!(f, "{}", self)
183    }
184}
185
186impl<D: DataRef> fmt::Display for GLWESwitchingKey<D> {
187    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
188        write!(
189            f,
190            "(GLWESwitchingKey: sk_in_n={} sk_out_n={}) {}",
191            self.sk_in_n, self.sk_out_n, self.key.data
192        )
193    }
194}
195
196impl<D: DataMut> FillUniform for GLWESwitchingKey<D> {
197    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
198        self.key.fill_uniform(source);
199    }
200}
201
202impl<D: DataMut> Reset for GLWESwitchingKey<D>
203where
204    MatZnx<D>: Reset,
205{
206    fn reset(&mut self) {
207        self.key.reset();
208        self.sk_in_n = 0;
209        self.sk_out_n = 0;
210    }
211}
212
213impl GLWESwitchingKey<Vec<u8>> {
214    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize, rank_out: usize) -> Self {
215        GLWESwitchingKey {
216            key: GGLWECiphertext::alloc(n, basek, k, rows, digits, rank_in, rank_out),
217            sk_in_n: 0,
218            sk_out_n: 0,
219        }
220    }
221
222    pub fn bytes_of(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize, rank_out: usize) -> usize {
223        GGLWECiphertext::<Vec<u8>>::bytes_of(n, basek, k, rows, digits, rank_in, rank_out)
224    }
225}
226
227impl<D: Data> Infos for GLWESwitchingKey<D> {
228    type Inner = MatZnx<D>;
229
230    fn inner(&self) -> &Self::Inner {
231        self.key.inner()
232    }
233
234    fn basek(&self) -> usize {
235        self.key.basek()
236    }
237
238    fn k(&self) -> usize {
239        self.key.k()
240    }
241}
242
243impl<D: Data> GLWESwitchingKey<D> {
244    pub fn rank(&self) -> usize {
245        self.key.data.cols_out() - 1
246    }
247
248    pub fn rank_in(&self) -> usize {
249        self.key.data.cols_in()
250    }
251
252    pub fn rank_out(&self) -> usize {
253        self.key.data.cols_out() - 1
254    }
255
256    pub fn digits(&self) -> usize {
257        self.key.digits()
258    }
259
260    pub fn sk_degree_in(&self) -> usize {
261        self.sk_in_n
262    }
263
264    pub fn sk_degree_out(&self) -> usize {
265        self.sk_out_n
266    }
267}
268
269impl<D: DataRef> GLWESwitchingKey<D> {
270    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
271        self.key.at(row, col)
272    }
273}
274
275impl<D: DataMut> GLWESwitchingKey<D> {
276    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
277        self.key.at_mut(row, col)
278    }
279}
280
281impl<D: DataMut> ReaderFrom for GLWESwitchingKey<D> {
282    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
283        self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
284        self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
285        self.key.read_from(reader)
286    }
287}
288
289impl<D: DataRef> WriterTo for GLWESwitchingKey<D> {
290    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
291        writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
292        writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
293        self.key.write_to(writer)
294    }
295}
296
297#[derive(PartialEq, Eq, Clone)]
298pub struct AutomorphismKey<D: Data> {
299    pub(crate) key: GLWESwitchingKey<D>,
300    pub(crate) p: i64,
301}
302
303impl<D: DataRef> fmt::Debug for AutomorphismKey<D> {
304    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
305        write!(f, "{}", self)
306    }
307}
308
309impl<D: DataMut> FillUniform for AutomorphismKey<D> {
310    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
311        self.key.fill_uniform(source);
312    }
313}
314
315impl<D: DataMut> Reset for AutomorphismKey<D>
316where
317    MatZnx<D>: Reset,
318{
319    fn reset(&mut self) {
320        self.key.reset();
321        self.p = 0;
322    }
323}
324
325impl<D: DataRef> fmt::Display for AutomorphismKey<D> {
326    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
327        write!(f, "(AutomorphismKey: p={}) {}", self.p, self.key)
328    }
329}
330
331impl AutomorphismKey<Vec<u8>> {
332    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self {
333        AutomorphismKey {
334            key: GLWESwitchingKey::alloc(n, basek, k, rows, digits, rank, rank),
335            p: 0,
336        }
337    }
338
339    pub fn bytes_of(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize {
340        GLWESwitchingKey::bytes_of(n, basek, k, rows, digits, rank, rank)
341    }
342}
343
344impl<D: Data> Infos for AutomorphismKey<D> {
345    type Inner = MatZnx<D>;
346
347    fn inner(&self) -> &Self::Inner {
348        &self.key.inner()
349    }
350
351    fn basek(&self) -> usize {
352        self.key.basek()
353    }
354
355    fn k(&self) -> usize {
356        self.key.k()
357    }
358}
359
360impl<D: Data> AutomorphismKey<D> {
361    pub fn p(&self) -> i64 {
362        self.p
363    }
364
365    pub fn digits(&self) -> usize {
366        self.key.digits()
367    }
368
369    pub fn rank(&self) -> usize {
370        self.key.rank()
371    }
372
373    pub fn rank_in(&self) -> usize {
374        self.key.rank_in()
375    }
376
377    pub fn rank_out(&self) -> usize {
378        self.key.rank_out()
379    }
380}
381
382impl<D: DataRef> AutomorphismKey<D> {
383    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
384        self.key.at(row, col)
385    }
386}
387
388impl<D: DataMut> AutomorphismKey<D> {
389    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
390        self.key.at_mut(row, col)
391    }
392}
393
394impl<D: DataMut> ReaderFrom for AutomorphismKey<D> {
395    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
396        self.p = reader.read_u64::<LittleEndian>()? as i64;
397        self.key.read_from(reader)
398    }
399}
400
401impl<D: DataRef> WriterTo for AutomorphismKey<D> {
402    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
403        writer.write_u64::<LittleEndian>(self.p as u64)?;
404        self.key.write_to(writer)
405    }
406}
407
408#[derive(PartialEq, Eq, Clone)]
409pub struct GLWETensorKey<D: Data> {
410    pub(crate) keys: Vec<GLWESwitchingKey<D>>,
411}
412
413impl<D: DataRef> fmt::Debug for GLWETensorKey<D> {
414    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
415        write!(f, "{}", self)
416    }
417}
418
419impl<D: DataMut> FillUniform for GLWETensorKey<D> {
420    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
421        self.keys
422            .iter_mut()
423            .for_each(|key: &mut GLWESwitchingKey<D>| key.fill_uniform(source))
424    }
425}
426
427impl<D: DataMut> Reset for GLWETensorKey<D>
428where
429    MatZnx<D>: Reset,
430{
431    fn reset(&mut self) {
432        self.keys
433            .iter_mut()
434            .for_each(|key: &mut GLWESwitchingKey<D>| key.reset())
435    }
436}
437
438impl<D: DataRef> fmt::Display for GLWETensorKey<D> {
439    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
440        writeln!(f, "(GLWETensorKey)",)?;
441        for (i, key) in self.keys.iter().enumerate() {
442            write!(f, "{}: {}", i, key)?;
443        }
444        Ok(())
445    }
446}
447
448impl GLWETensorKey<Vec<u8>> {
449    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self {
450        let mut keys: Vec<GLWESwitchingKey<Vec<u8>>> = Vec::new();
451        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
452        (0..pairs).for_each(|_| {
453            keys.push(GLWESwitchingKey::alloc(n, basek, k, rows, digits, 1, rank));
454        });
455        Self { keys: keys }
456    }
457
458    pub fn bytes_of(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> usize {
459        let pairs: usize = (((rank + 1) * rank) >> 1).max(1);
460        pairs * GLWESwitchingKey::<Vec<u8>>::bytes_of(n, basek, k, rows, digits, 1, rank)
461    }
462}
463
464impl<D: Data> Infos for GLWETensorKey<D> {
465    type Inner = MatZnx<D>;
466
467    fn inner(&self) -> &Self::Inner {
468        &self.keys[0].inner()
469    }
470
471    fn basek(&self) -> usize {
472        self.keys[0].basek()
473    }
474
475    fn k(&self) -> usize {
476        self.keys[0].k()
477    }
478}
479
480impl<D: Data> GLWETensorKey<D> {
481    pub fn rank(&self) -> usize {
482        self.keys[0].rank()
483    }
484
485    pub fn rank_in(&self) -> usize {
486        self.keys[0].rank_in()
487    }
488
489    pub fn rank_out(&self) -> usize {
490        self.keys[0].rank_out()
491    }
492
493    pub fn digits(&self) -> usize {
494        self.keys[0].digits()
495    }
496}
497
498impl<D: DataMut> GLWETensorKey<D> {
499    // Returns a mutable reference to GLWESwitchingKey_{s}(s[i] * s[j])
500    pub fn at_mut(&mut self, mut i: usize, mut j: usize) -> &mut GLWESwitchingKey<D> {
501        if i > j {
502            std::mem::swap(&mut i, &mut j);
503        };
504        let rank: usize = self.rank();
505        &mut self.keys[i * rank + j - (i * (i + 1) / 2)]
506    }
507}
508
509impl<D: DataRef> GLWETensorKey<D> {
510    // Returns a reference to GLWESwitchingKey_{s}(s[i] * s[j])
511    pub fn at(&self, mut i: usize, mut j: usize) -> &GLWESwitchingKey<D> {
512        if i > j {
513            std::mem::swap(&mut i, &mut j);
514        };
515        let rank: usize = self.rank();
516        &self.keys[i * rank + j - (i * (i + 1) / 2)]
517    }
518}
519
520impl<D: DataMut> ReaderFrom for GLWETensorKey<D> {
521    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
522        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
523        if self.keys.len() != len {
524            return Err(std::io::Error::new(
525                std::io::ErrorKind::InvalidData,
526                format!("self.keys.len()={} != read len={}", self.keys.len(), len),
527            ));
528        }
529        for key in &mut self.keys {
530            key.read_from(reader)?;
531        }
532        Ok(())
533    }
534}
535
536impl<D: DataRef> WriterTo for GLWETensorKey<D> {
537    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
538        writer.write_u64::<LittleEndian>(self.keys.len() as u64)?;
539        for key in &self.keys {
540            key.write_to(writer)?;
541        }
542        Ok(())
543    }
544}