1use poulpy_backend::hal::{
  2    api::{FillUniform, Reset},
  3    layouts::{Data, DataMut, DataRef, MatZnx, ReaderFrom, WriterTo},
  4    source::Source,
  5};
  6
  7use crate::layouts::{GGLWECiphertext, GLWECiphertext, Infos};
  8use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
  9
 10use std::fmt;
 11
 12#[derive(PartialEq, Eq, Clone)]
 13pub struct GGLWESwitchingKey<D: Data> {
 14    pub(crate) key: GGLWECiphertext<D>,
 15    pub(crate) sk_in_n: usize,  // Degree of sk_in
 16    pub(crate) sk_out_n: usize, // Degree of sk_out
 17}
 18
 19impl<D: DataRef> fmt::Debug for GGLWESwitchingKey<D> {
 20    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 21        write!(f, "{}", self)
 22    }
 23}
 24
 25impl<D: DataRef> fmt::Display for GGLWESwitchingKey<D> {
 26    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 27        write!(
 28            f,
 29            "(GLWESwitchingKey: sk_in_n={} sk_out_n={}) {}",
 30            self.sk_in_n, self.sk_out_n, self.key.data
 31        )
 32    }
 33}
 34
 35impl<D: DataMut> FillUniform for GGLWESwitchingKey<D> {
 36    fn fill_uniform(&mut self, source: &mut Source) {
 37        self.key.fill_uniform(source);
 38    }
 39}
 40
 41impl<D: DataMut> Reset for GGLWESwitchingKey<D>
 42where
 43    MatZnx<D>: Reset,
 44{
 45    fn reset(&mut self) {
 46        self.key.reset();
 47        self.sk_in_n = 0;
 48        self.sk_out_n = 0;
 49    }
 50}
 51
 52impl GGLWESwitchingKey<Vec<u8>> {
 53    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize, rank_out: usize) -> Self {
 54        GGLWESwitchingKey {
 55            key: GGLWECiphertext::alloc(n, basek, k, rows, digits, rank_in, rank_out),
 56            sk_in_n: 0,
 57            sk_out_n: 0,
 58        }
 59    }
 60
 61    pub fn bytes_of(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize, rank_out: usize) -> usize {
 62        GGLWECiphertext::<Vec<u8>>::bytes_of(n, basek, k, rows, digits, rank_in, rank_out)
 63    }
 64}
 65
 66impl<D: Data> Infos for GGLWESwitchingKey<D> {
 67    type Inner = MatZnx<D>;
 68
 69    fn inner(&self) -> &Self::Inner {
 70        self.key.inner()
 71    }
 72
 73    fn basek(&self) -> usize {
 74        self.key.basek()
 75    }
 76
 77    fn k(&self) -> usize {
 78        self.key.k()
 79    }
 80}
 81
 82impl<D: Data> GGLWESwitchingKey<D> {
 83    pub fn rank(&self) -> usize {
 84        self.key.data.cols_out() - 1
 85    }
 86
 87    pub fn rank_in(&self) -> usize {
 88        self.key.data.cols_in()
 89    }
 90
 91    pub fn rank_out(&self) -> usize {
 92        self.key.data.cols_out() - 1
 93    }
 94
 95    pub fn digits(&self) -> usize {
 96        self.key.digits()
 97    }
 98
 99    pub fn sk_degree_in(&self) -> usize {
100        self.sk_in_n
101    }
102
103    pub fn sk_degree_out(&self) -> usize {
104        self.sk_out_n
105    }
106}
107
108impl<D: DataRef> GGLWESwitchingKey<D> {
109    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
110        self.key.at(row, col)
111    }
112}
113
114impl<D: DataMut> GGLWESwitchingKey<D> {
115    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
116        self.key.at_mut(row, col)
117    }
118}
119
120impl<D: DataMut> ReaderFrom for GGLWESwitchingKey<D> {
121    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
122        self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
123        self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
124        self.key.read_from(reader)
125    }
126}
127
128impl<D: DataRef> WriterTo for GGLWESwitchingKey<D> {
129    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
130        writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
131        writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
132        self.key.write_to(writer)
133    }
134}