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