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