1use poulpy_hal::{
  2    layouts::{Backend, Data, DataMut, DataRef, FillUniform, Module, ReaderFrom, WriterTo},
  3    source::Source,
  4};
  5
  6use crate::layouts::{
  7    AutomorphismKey, Base2K, Degree, Dnum, Dsize, GGLWECompressed, GGLWECompressedSeedMut, GGLWECompressedToMut,
  8    GGLWECompressedToRef, GGLWEDecompress, GGLWEInfos, GGLWEToMut, GLWECompressed, GLWECompressedToMut, GLWECompressedToRef,
  9    GLWEDecompress, GLWEInfos, LWEInfos, Rank, TorusPrecision,
 10    prepared::{GetAutomorphismGaloisElement, SetAutomorphismGaloisElement},
 11};
 12use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 13use std::fmt;
 14
 15#[derive(PartialEq, Eq, Clone)]
 16pub struct AutomorphismKeyCompressed<D: Data> {
 17    pub(crate) key: GGLWECompressed<D>,
 18    pub(crate) p: i64,
 19}
 20
 21impl<D: Data> GetAutomorphismGaloisElement for AutomorphismKeyCompressed<D> {
 22    fn p(&self) -> i64 {
 23        self.p
 24    }
 25}
 26
 27impl<D: Data> LWEInfos for AutomorphismKeyCompressed<D> {
 28    fn n(&self) -> Degree {
 29        self.key.n()
 30    }
 31
 32    fn base2k(&self) -> Base2K {
 33        self.key.base2k()
 34    }
 35
 36    fn k(&self) -> TorusPrecision {
 37        self.key.k()
 38    }
 39
 40    fn size(&self) -> usize {
 41        self.key.size()
 42    }
 43}
 44
 45impl<D: Data> GLWEInfos for AutomorphismKeyCompressed<D> {
 46    fn rank(&self) -> Rank {
 47        self.rank_out()
 48    }
 49}
 50
 51impl<D: Data> GGLWEInfos for AutomorphismKeyCompressed<D> {
 52    fn rank_in(&self) -> Rank {
 53        self.key.rank_in()
 54    }
 55
 56    fn rank_out(&self) -> Rank {
 57        self.key.rank_out()
 58    }
 59
 60    fn dsize(&self) -> Dsize {
 61        self.key.dsize()
 62    }
 63
 64    fn dnum(&self) -> Dnum {
 65        self.key.dnum()
 66    }
 67}
 68
 69impl<D: DataRef> fmt::Debug for AutomorphismKeyCompressed<D> {
 70    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 71        write!(f, "{self}")
 72    }
 73}
 74
 75impl<D: DataMut> FillUniform for AutomorphismKeyCompressed<D> {
 76    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
 77        self.key.fill_uniform(log_bound, source);
 78    }
 79}
 80
 81impl<D: DataRef> fmt::Display for AutomorphismKeyCompressed<D> {
 82    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 83        write!(f, "(AutomorphismKeyCompressed: p={}) {}", self.p, self.key)
 84    }
 85}
 86
 87impl AutomorphismKeyCompressed<Vec<u8>> {
 88    pub fn alloc_from_infos<A>(infos: &A) -> Self
 89    where
 90        A: GGLWEInfos,
 91    {
 92        Self::alloc(
 93            infos.n(),
 94            infos.base2k(),
 95            infos.k(),
 96            infos.rank(),
 97            infos.dnum(),
 98            infos.dsize(),
 99        )
100    }
101
102    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self {
103        AutomorphismKeyCompressed {
104            key: GGLWECompressed::alloc(n, base2k, k, rank, rank, dnum, dsize),
105            p: 0,
106        }
107    }
108
109    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
110    where
111        A: GGLWEInfos,
112    {
113        Self::bytes_of(
114            infos.n(),
115            infos.base2k(),
116            infos.k(),
117            infos.rank(),
118            infos.dnum(),
119            infos.dsize(),
120        )
121    }
122
123    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
124        GGLWECompressed::bytes_of(n, base2k, k, rank, dnum, dsize)
125    }
126}
127
128impl<D: DataMut> ReaderFrom for AutomorphismKeyCompressed<D> {
129    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
130        self.p = reader.read_u64::<LittleEndian>()? as i64;
131        self.key.read_from(reader)
132    }
133}
134
135impl<D: DataRef> WriterTo for AutomorphismKeyCompressed<D> {
136    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
137        writer.write_u64::<LittleEndian>(self.p as u64)?;
138        self.key.write_to(writer)
139    }
140}
141
142pub trait AutomorphismKeyDecompress
143where
144    Self: GGLWEDecompress,
145{
146    fn decompress_automorphism_key<R, O>(&self, res: &mut R, other: &O)
147    where
148        R: GGLWEToMut + SetAutomorphismGaloisElement,
149        O: GGLWECompressedToRef + GetAutomorphismGaloisElement,
150    {
151        self.decompress_gglwe(res, other);
152        res.set_p(other.p());
153    }
154}
155
156impl<B: Backend> AutomorphismKeyDecompress for Module<B> where Self: GLWEDecompress {}
157
158impl<D: DataMut> AutomorphismKey<D>
159where
160    Self: SetAutomorphismGaloisElement,
161{
162    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
163    where
164        O: GGLWECompressedToRef + GetAutomorphismGaloisElement,
165        M: AutomorphismKeyDecompress,
166    {
167        module.decompress_automorphism_key(self, other);
168    }
169}
170
171impl<D: DataRef> GGLWECompressedToRef for AutomorphismKeyCompressed<D> {
172    fn to_ref(&self) -> GGLWECompressed<&[u8]> {
173        self.key.to_ref()
174    }
175}
176
177impl<D: DataMut> GGLWECompressedToMut for AutomorphismKeyCompressed<D> {
178    fn to_mut(&mut self) -> GGLWECompressed<&mut [u8]> {
179        self.key.to_mut()
180    }
181}
182
183pub trait AutomorphismKeyCompressedToRef {
184    fn to_ref(&self) -> AutomorphismKeyCompressed<&[u8]>;
185}
186
187impl<D: DataRef> AutomorphismKeyCompressedToRef for AutomorphismKeyCompressed<D>
188where
189    GLWECompressed<D>: GLWECompressedToRef,
190{
191    fn to_ref(&self) -> AutomorphismKeyCompressed<&[u8]> {
192        AutomorphismKeyCompressed {
193            key: self.key.to_ref(),
194            p: self.p,
195        }
196    }
197}
198
199pub trait AutomorphismKeyCompressedToMut {
200    fn to_mut(&mut self) -> AutomorphismKeyCompressed<&mut [u8]>;
201}
202
203impl<D: DataMut> AutomorphismKeyCompressedToMut for AutomorphismKeyCompressed<D>
204where
205    GLWECompressed<D>: GLWECompressedToMut,
206{
207    fn to_mut(&mut self) -> AutomorphismKeyCompressed<&mut [u8]> {
208        AutomorphismKeyCompressed {
209            p: self.p,
210            key: self.key.to_mut(),
211        }
212    }
213}
214
215impl<D: DataMut> GGLWECompressedSeedMut for AutomorphismKeyCompressed<D> {
216    fn seed_mut(&mut self) -> &mut Vec<[u8; 32]> {
217        &mut self.key.seed
218    }
219}
220
221impl<D: DataMut> SetAutomorphismGaloisElement for AutomorphismKeyCompressed<D> {
222    fn set_p(&mut self, p: i64) {
223        self.p = p
224    }
225}