1use backend::hal::{
  2    api::{FillUniform, Reset, VecZnxCopy, VecZnxFillUniform},
  3    layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, WriterTo},
  4};
  5
  6use crate::layouts::{GGLWEAutomorphismKey, Infos, compressed::GGLWESwitchingKeyCompressed};
  7use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
  8use std::fmt;
  9
 10use crate::trait_families::Decompress;
 11
 12#[derive(PartialEq, Eq, Clone)]
 13pub struct GGLWEAutomorphismKeyCompressed<D: Data> {
 14    pub(crate) key: GGLWESwitchingKeyCompressed<D>,
 15    pub(crate) p: i64,
 16}
 17
 18impl<D: DataRef> fmt::Debug for GGLWEAutomorphismKeyCompressed<D> {
 19    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 20        write!(f, "{}", self)
 21    }
 22}
 23
 24impl<D: DataMut> FillUniform for GGLWEAutomorphismKeyCompressed<D> {
 25    fn fill_uniform(&mut self, source: &mut sampling::source::Source) {
 26        self.key.fill_uniform(source);
 27    }
 28}
 29
 30impl<D: DataMut> Reset for GGLWEAutomorphismKeyCompressed<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 GGLWEAutomorphismKeyCompressed<D> {
 41    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 42        write!(f, "(AutomorphismKeyCompressed: p={}) {}", self.p, self.key)
 43    }
 44}
 45
 46impl GGLWEAutomorphismKeyCompressed<Vec<u8>> {
 47    pub fn alloc(n: usize, basek: usize, k: usize, rows: usize, digits: usize, rank: usize) -> Self {
 48        GGLWEAutomorphismKeyCompressed {
 49            key: GGLWESwitchingKeyCompressed::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        GGLWESwitchingKeyCompressed::<Vec<u8>>::bytes_of(n, basek, k, rows, digits, rank)
 56    }
 57}
 58
 59impl<D: Data> Infos for GGLWEAutomorphismKeyCompressed<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> GGLWEAutomorphismKeyCompressed<D> {
 76    pub fn rank(&self) -> usize {
 77        self.key.rank()
 78    }
 79
 80    pub fn digits(&self) -> usize {
 81        self.key.digits()
 82    }
 83
 84    pub fn rank_in(&self) -> usize {
 85        self.key.rank_in()
 86    }
 87
 88    pub fn rank_out(&self) -> usize {
 89        self.key.rank_out()
 90    }
 91}
 92
 93impl<D: DataMut> ReaderFrom for GGLWEAutomorphismKeyCompressed<D> {
 94    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
 95        self.p = reader.read_u64::<LittleEndian>()? as i64;
 96        self.key.read_from(reader)
 97    }
 98}
 99
100impl<D: DataRef> WriterTo for GGLWEAutomorphismKeyCompressed<D> {
101    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
102        writer.write_u64::<LittleEndian>(self.p as u64)?;
103        self.key.write_to(writer)
104    }
105}
106
107impl<D: DataMut, DR: DataRef, B: Backend> Decompress<B, GGLWEAutomorphismKeyCompressed<DR>> for GGLWEAutomorphismKey<D> {
108    fn decompress(&mut self, module: &Module<B>, other: &GGLWEAutomorphismKeyCompressed<DR>)
109    where
110        Module<B>: VecZnxFillUniform + VecZnxCopy,
111    {
112        self.key.decompress(module, &other.key);
113        self.p = other.p;
114    }
115}