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