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