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