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