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