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, GGLWESwitchingKey, GGLWESwitchingKeyToMut, GLWECiphertext, GLWEInfos, LWEInfos,
  8    Rank, TorusPrecision,
  9};
 10use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 11
 12use std::fmt;
 13
 14#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 15pub struct GGLWEAutomorphismKeyLayout {
 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 GGLWEAutomorphismKey<D: Data> {
 26    pub(crate) key: GGLWESwitchingKey<D>,
 27    pub(crate) p: i64,
 28}
 29
 30impl<D: Data> GGLWEAutomorphismKey<D> {
 31    pub fn p(&self) -> i64 {
 32        self.p
 33    }
 34}
 35
 36impl<D: Data> LWEInfos for GGLWEAutomorphismKey<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 GGLWEAutomorphismKey<D> {
 55    fn rank(&self) -> Rank {
 56        self.rank_out()
 57    }
 58}
 59
 60impl<D: Data> GGLWEInfos for GGLWEAutomorphismKey<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 GGLWEAutomorphismKeyLayout {
 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 GGLWEAutomorphismKeyLayout {
 93    fn rank(&self) -> Rank {
 94        self.rank
 95    }
 96}
 97
 98impl GGLWEInfos for GGLWEAutomorphismKeyLayout {
 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 GGLWEAutomorphismKey<D> {
117    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
118        write!(f, "{self}")
119    }
120}
121
122impl<D: DataMut> FillUniform for GGLWEAutomorphismKey<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 GGLWEAutomorphismKey<D> {
129    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
130        write!(f, "(AutomorphismKey: p={}) {}", self.p, self.key)
131    }
132}
133
134impl GGLWEAutomorphismKey<Vec<u8>> {
135    pub fn alloc<A>(infos: &A) -> Self
136    where
137        A: GGLWEInfos,
138    {
139        assert_eq!(
140            infos.rank_in(),
141            infos.rank_out(),
142            "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
143        );
144        GGLWEAutomorphismKey {
145            key: GGLWESwitchingKey::alloc(infos),
146            p: 0,
147        }
148    }
149
150    pub fn alloc_with(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self {
151        GGLWEAutomorphismKey {
152            key: GGLWESwitchingKey::alloc_with(n, base2k, k, rank, rank, dnum, dsize),
153            p: 0,
154        }
155    }
156
157    pub fn alloc_bytes<A>(infos: &A) -> usize
158    where
159        A: GGLWEInfos,
160    {
161        assert_eq!(
162            infos.rank_in(),
163            infos.rank_out(),
164            "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
165        );
166        GGLWESwitchingKey::alloc_bytes(infos)
167    }
168
169    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
170        GGLWESwitchingKey::alloc_bytes_with(n, base2k, k, rank, rank, dnum, dsize)
171    }
172}
173
174pub trait GGLWEAutomorphismKeyToMut {
175    fn to_mut(&mut self) -> GGLWEAutomorphismKey<&mut [u8]>;
176}
177
178impl<D: DataMut> GGLWEAutomorphismKeyToMut for GGLWEAutomorphismKey<D>
179where
180    GGLWESwitchingKey<D>: GGLWESwitchingKeyToMut,
181{
182    fn to_mut(&mut self) -> GGLWEAutomorphismKey<&mut [u8]> {
183        GGLWEAutomorphismKey {
184            key: self.key.to_mut(),
185            p: self.p,
186        }
187    }
188}
189
190pub trait GGLWEAutomorphismKeyToRef {
191    fn to_ref(&self) -> GGLWEAutomorphismKey<&[u8]>;
192}
193
194impl<D: DataRef> GGLWEAutomorphismKey<D> {
195    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
196        self.key.at(row, col)
197    }
198}
199
200impl<D: DataMut> GGLWEAutomorphismKey<D> {
201    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
202        self.key.at_mut(row, col)
203    }
204}
205
206impl<D: DataMut> ReaderFrom for GGLWEAutomorphismKey<D> {
207    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
208        self.p = reader.read_u64::<LittleEndian>()? as i64;
209        self.key.read_from(reader)
210    }
211}
212
213impl<D: DataRef> WriterTo for GGLWEAutomorphismKey<D> {
214    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
215        writer.write_u64::<LittleEndian>(self.p as u64)?;
216        self.key.write_to(writer)
217    }
218}