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