1use poulpy_hal::{
  2    layouts::{Data, DataMut, DataRef, FillUniform, ReaderFrom, WriterTo},
  3    source::Source,
  4};
  5
  6use crate::layouts::{
  7    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWE, GLWEInfos, GLWESwitchingKey, GLWESwitchingKeyToMut, GLWESwitchingKeyToRef,
  8    LWEInfos, Rank, 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: Degree,
 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) -> Degree {
 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) -> Degree {
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 AutomorphismKey<Vec<u8>> {
148    pub fn alloc_from_infos<A>(infos: &A) -> Self
149    where
150        A: GGLWEInfos,
151    {
152        Self::alloc(
153            infos.n(),
154            infos.base2k(),
155            infos.k(),
156            infos.rank(),
157            infos.dnum(),
158            infos.dsize(),
159        )
160    }
161
162    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self {
163        AutomorphismKey {
164            key: GLWESwitchingKey::alloc(n, base2k, k, rank, rank, dnum, dsize),
165            p: 0,
166        }
167    }
168
169    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
170    where
171        A: GGLWEInfos,
172    {
173        assert_eq!(
174            infos.rank_in(),
175            infos.rank_out(),
176            "rank_in != rank_out is not supported for AutomorphismKey"
177        );
178        Self::bytes_of(
179            infos.n(),
180            infos.base2k(),
181            infos.k(),
182            infos.rank(),
183            infos.dnum(),
184            infos.dsize(),
185        )
186    }
187
188    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
189        GLWESwitchingKey::bytes_of(n, base2k, k, rank, rank, dnum, dsize)
190    }
191}
192
193pub trait AutomorphismKeyToMut {
194    fn to_mut(&mut self) -> AutomorphismKey<&mut [u8]>;
195}
196
197impl<D: DataMut> AutomorphismKeyToMut for AutomorphismKey<D>
198where
199    GLWESwitchingKey<D>: GLWESwitchingKeyToMut,
200{
201    fn to_mut(&mut self) -> AutomorphismKey<&mut [u8]> {
202        AutomorphismKey {
203            key: self.key.to_mut(),
204            p: self.p,
205        }
206    }
207}
208
209pub trait AutomorphismKeyToRef {
210    fn to_ref(&self) -> AutomorphismKey<&[u8]>;
211}
212
213impl<D: DataRef> AutomorphismKeyToRef for AutomorphismKey<D>
214where
215    GLWESwitchingKey<D>: GLWESwitchingKeyToRef,
216{
217    fn to_ref(&self) -> AutomorphismKey<&[u8]> {
218        AutomorphismKey {
219            p: self.p,
220            key: self.key.to_ref(),
221        }
222    }
223}
224
225impl<D: DataRef> AutomorphismKey<D> {
226    pub fn at(&self, row: usize, col: usize) -> GLWE<&[u8]> {
227        self.key.at(row, col)
228    }
229}
230
231impl<D: DataMut> AutomorphismKey<D> {
232    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWE<&mut [u8]> {
233        self.key.at_mut(row, col)
234    }
235}
236
237impl<D: DataMut> ReaderFrom for AutomorphismKey<D> {
238    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
239        self.p = reader.read_u64::<LittleEndian>()? as i64;
240        self.key.read_from(reader)
241    }
242}
243
244impl<D: DataRef> WriterTo for AutomorphismKey<D> {
245    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
246        writer.write_u64::<LittleEndian>(self.p as u64)?;
247        self.key.write_to(writer)
248    }
249}