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};
 10use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 11
 12use std::fmt;
 13
 14#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 15pub struct AutomorphismKeyLayout {
 16    pub n: RingDegree,
 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) -> RingDegree {
 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) -> RingDegree {
 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, M>(module: &M, infos: &A) -> Self
192    where
193        A: GGLWEInfos,
194        M: AutomorphismKeyAlloc,
195    {
196        module.alloc_automorphism_key_from_infos(infos)
197    }
198
199    pub fn alloc_with<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self
200    where
201        M: AutomorphismKeyAlloc,
202    {
203        module.alloc_automorphism_key(base2k, k, rank, dnum, dsize)
204    }
205
206    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
207    where
208        A: GGLWEInfos,
209        M: AutomorphismKeyAlloc,
210    {
211        module.bytes_of_automorphism_key_from_infos(infos)
212    }
213
214    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize
215    where
216        M: AutomorphismKeyAlloc,
217    {
218        module.bytes_of_automorphism_key(base2k, k, rank, dnum, dsize)
219    }
220}
221
222pub trait AutomorphismKeyToMut {
223    fn to_mut(&mut self) -> AutomorphismKey<&mut [u8]>;
224}
225
226impl<D: DataMut> AutomorphismKeyToMut for AutomorphismKey<D>
227where
228    GLWESwitchingKey<D>: GLWESwitchingKeyToMut,
229{
230    fn to_mut(&mut self) -> AutomorphismKey<&mut [u8]> {
231        AutomorphismKey {
232            key: self.key.to_mut(),
233            p: self.p,
234        }
235    }
236}
237
238pub trait AutomorphismKeyToRef {
239    fn to_ref(&self) -> AutomorphismKey<&[u8]>;
240}
241
242impl<D: DataRef> AutomorphismKeyToRef for AutomorphismKey<D>
243where
244    GLWESwitchingKey<D>: GLWESwitchingKeyToRef,
245{
246    fn to_ref(&self) -> AutomorphismKey<&[u8]> {
247        AutomorphismKey {
248            p: self.p,
249            key: self.key.to_ref(),
250        }
251    }
252}
253
254impl<D: DataRef> AutomorphismKey<D> {
255    pub fn at(&self, row: usize, col: usize) -> GLWE<&[u8]> {
256        self.key.at(row, col)
257    }
258}
259
260impl<D: DataMut> AutomorphismKey<D> {
261    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWE<&mut [u8]> {
262        self.key.at_mut(row, col)
263    }
264}
265
266impl<D: DataMut> ReaderFrom for AutomorphismKey<D> {
267    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
268        self.p = reader.read_u64::<LittleEndian>()? as i64;
269        self.key.read_from(reader)
270    }
271}
272
273impl<D: DataRef> WriterTo for AutomorphismKey<D> {
274    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
275        writer.write_u64::<LittleEndian>(self.p as u64)?;
276        self.key.write_to(writer)
277    }
278}