1use poulpy_hal::{
  2    layouts::{Data, DataMut, DataRef, FillUniform, ReaderFrom, WriterTo},
  3    source::Source,
  4};
  5
  6use crate::layouts::{
  7    Base2K, Degree, Dnum, Dsize, GGLWE, GGLWEInfos, GGLWEToMut, GGLWEToRef, GLWE, GLWEInfos, LWEInfos, Rank, TorusPrecision,
  8    prepared::{GetAutomorphismGaloisElement, SetAutomorphismGaloisElement},
  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: GGLWE<D>,
 27    pub(crate) p: i64,
 28}
 29
 30impl<D: DataMut> SetAutomorphismGaloisElement for AutomorphismKey<D> {
 31    fn set_p(&mut self, p: i64) {
 32        self.p = p
 33    }
 34}
 35
 36impl<D: DataRef> GetAutomorphismGaloisElement for AutomorphismKey<D> {
 37    fn p(&self) -> i64 {
 38        self.p
 39    }
 40}
 41
 42impl<D: Data> AutomorphismKey<D> {
 43    pub fn p(&self) -> i64 {
 44        self.p
 45    }
 46}
 47
 48impl<D: Data> LWEInfos for AutomorphismKey<D> {
 49    fn n(&self) -> Degree {
 50        self.key.n()
 51    }
 52
 53    fn base2k(&self) -> Base2K {
 54        self.key.base2k()
 55    }
 56
 57    fn k(&self) -> TorusPrecision {
 58        self.key.k()
 59    }
 60
 61    fn size(&self) -> usize {
 62        self.key.size()
 63    }
 64}
 65
 66impl<D: Data> GLWEInfos for AutomorphismKey<D> {
 67    fn rank(&self) -> Rank {
 68        self.rank_out()
 69    }
 70}
 71
 72impl<D: Data> GGLWEInfos for AutomorphismKey<D> {
 73    fn rank_in(&self) -> Rank {
 74        self.key.rank_in()
 75    }
 76
 77    fn rank_out(&self) -> Rank {
 78        self.key.rank_out()
 79    }
 80
 81    fn dsize(&self) -> Dsize {
 82        self.key.dsize()
 83    }
 84
 85    fn dnum(&self) -> Dnum {
 86        self.key.dnum()
 87    }
 88}
 89
 90impl LWEInfos for AutomorphismKeyLayout {
 91    fn base2k(&self) -> Base2K {
 92        self.base2k
 93    }
 94
 95    fn k(&self) -> TorusPrecision {
 96        self.k
 97    }
 98
 99    fn n(&self) -> Degree {
100        self.n
101    }
102}
103
104impl GLWEInfos for AutomorphismKeyLayout {
105    fn rank(&self) -> Rank {
106        self.rank
107    }
108}
109
110impl GGLWEInfos for AutomorphismKeyLayout {
111    fn rank_in(&self) -> Rank {
112        self.rank
113    }
114
115    fn dsize(&self) -> Dsize {
116        self.dsize
117    }
118
119    fn rank_out(&self) -> Rank {
120        self.rank
121    }
122
123    fn dnum(&self) -> Dnum {
124        self.dnum
125    }
126}
127
128impl<D: DataRef> fmt::Debug for AutomorphismKey<D> {
129    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
130        write!(f, "{self}")
131    }
132}
133
134impl<D: DataMut> FillUniform for AutomorphismKey<D> {
135    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
136        self.key.fill_uniform(log_bound, source);
137    }
138}
139
140impl<D: DataRef> fmt::Display for AutomorphismKey<D> {
141    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
142        write!(f, "(AutomorphismKey: p={}) {}", self.p, self.key)
143    }
144}
145
146impl AutomorphismKey<Vec<u8>> {
147    pub fn alloc_from_infos<A>(infos: &A) -> Self
148    where
149        A: GGLWEInfos,
150    {
151        Self::alloc(
152            infos.n(),
153            infos.base2k(),
154            infos.k(),
155            infos.rank(),
156            infos.dnum(),
157            infos.dsize(),
158        )
159    }
160
161    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self {
162        AutomorphismKey {
163            key: GGLWE::alloc(n, base2k, k, rank, rank, dnum, dsize),
164            p: 0,
165        }
166    }
167
168    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
169    where
170        A: GGLWEInfos,
171    {
172        assert_eq!(
173            infos.rank_in(),
174            infos.rank_out(),
175            "rank_in != rank_out is not supported for AutomorphismKey"
176        );
177        Self::bytes_of(
178            infos.n(),
179            infos.base2k(),
180            infos.k(),
181            infos.rank(),
182            infos.dnum(),
183            infos.dsize(),
184        )
185    }
186
187    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
188        GGLWE::bytes_of(n, base2k, k, rank, rank, dnum, dsize)
189    }
190}
191
192impl<D: DataMut> GGLWEToMut for AutomorphismKey<D> {
193    fn to_mut(&mut self) -> GGLWE<&mut [u8]> {
194        self.key.to_mut()
195    }
196}
197
198impl<D: DataMut> GGLWEToRef for AutomorphismKey<D> {
199    fn to_ref(&self) -> GGLWE<&[u8]> {
200        self.key.to_ref()
201    }
202}
203
204impl<D: DataRef> AutomorphismKey<D> {
205    pub fn at(&self, row: usize, col: usize) -> GLWE<&[u8]> {
206        self.key.at(row, col)
207    }
208}
209
210impl<D: DataMut> AutomorphismKey<D> {
211    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWE<&mut [u8]> {
212        self.key.at_mut(row, col)
213    }
214}
215
216impl<D: DataMut> ReaderFrom for AutomorphismKey<D> {
217    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
218        self.p = reader.read_u64::<LittleEndian>()? as i64;
219        self.key.read_from(reader)
220    }
221}
222
223impl<D: DataRef> WriterTo for AutomorphismKey<D> {
224    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
225        writer.write_u64::<LittleEndian>(self.p as u64)?;
226        self.key.write_to(writer)
227    }
228}