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, GLWECiphertext, GLWEInfos, LWEInfos, Rank, TorusPrecision,
  8};
  9use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 10
 11use std::fmt;
 12
 13#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 14pub struct GGLWEAutomorphismKeyLayout {
 15    pub n: Degree,
 16    pub base2k: Base2K,
 17    pub k: TorusPrecision,
 18    pub rank: Rank,
 19    pub dnum: Dnum,
 20    pub dsize: Dsize,
 21}
 22
 23#[derive(PartialEq, Eq, Clone)]
 24pub struct GGLWEAutomorphismKey<D: Data> {
 25    pub(crate) key: GGLWESwitchingKey<D>,
 26    pub(crate) p: i64,
 27}
 28
 29impl<D: Data> GGLWEAutomorphismKey<D> {
 30    pub fn p(&self) -> i64 {
 31        self.p
 32    }
 33}
 34
 35impl<D: Data> LWEInfos for GGLWEAutomorphismKey<D> {
 36    fn n(&self) -> Degree {
 37        self.key.n()
 38    }
 39
 40    fn base2k(&self) -> Base2K {
 41        self.key.base2k()
 42    }
 43
 44    fn k(&self) -> TorusPrecision {
 45        self.key.k()
 46    }
 47
 48    fn size(&self) -> usize {
 49        self.key.size()
 50    }
 51}
 52
 53impl<D: Data> GLWEInfos for GGLWEAutomorphismKey<D> {
 54    fn rank(&self) -> Rank {
 55        self.rank_out()
 56    }
 57}
 58
 59impl<D: Data> GGLWEInfos for GGLWEAutomorphismKey<D> {
 60    fn rank_in(&self) -> Rank {
 61        self.key.rank_in()
 62    }
 63
 64    fn rank_out(&self) -> Rank {
 65        self.key.rank_out()
 66    }
 67
 68    fn dsize(&self) -> Dsize {
 69        self.key.dsize()
 70    }
 71
 72    fn dnum(&self) -> Dnum {
 73        self.key.dnum()
 74    }
 75}
 76
 77impl LWEInfos for GGLWEAutomorphismKeyLayout {
 78    fn base2k(&self) -> Base2K {
 79        self.base2k
 80    }
 81
 82    fn k(&self) -> TorusPrecision {
 83        self.k
 84    }
 85
 86    fn n(&self) -> Degree {
 87        self.n
 88    }
 89}
 90
 91impl GLWEInfos for GGLWEAutomorphismKeyLayout {
 92    fn rank(&self) -> Rank {
 93        self.rank
 94    }
 95}
 96
 97impl GGLWEInfos for GGLWEAutomorphismKeyLayout {
 98    fn rank_in(&self) -> Rank {
 99        self.rank
100    }
101
102    fn dsize(&self) -> Dsize {
103        self.dsize
104    }
105
106    fn rank_out(&self) -> Rank {
107        self.rank
108    }
109
110    fn dnum(&self) -> Dnum {
111        self.dnum
112    }
113}
114
115impl<D: DataRef> fmt::Debug for GGLWEAutomorphismKey<D> {
116    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
117        write!(f, "{self}")
118    }
119}
120
121impl<D: DataMut> FillUniform for GGLWEAutomorphismKey<D> {
122    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
123        self.key.fill_uniform(log_bound, source);
124    }
125}
126
127impl<D: DataRef> fmt::Display for GGLWEAutomorphismKey<D> {
128    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
129        write!(f, "(AutomorphismKey: p={}) {}", self.p, self.key)
130    }
131}
132
133impl GGLWEAutomorphismKey<Vec<u8>> {
134    pub fn alloc<A>(infos: &A) -> Self
135    where
136        A: GGLWEInfos,
137    {
138        assert_eq!(
139            infos.rank_in(),
140            infos.rank_out(),
141            "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
142        );
143        GGLWEAutomorphismKey {
144            key: GGLWESwitchingKey::alloc(infos),
145            p: 0,
146        }
147    }
148
149    pub fn alloc_with(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> Self {
150        GGLWEAutomorphismKey {
151            key: GGLWESwitchingKey::alloc_with(n, base2k, k, rank, rank, dnum, dsize),
152            p: 0,
153        }
154    }
155
156    pub fn alloc_bytes<A>(infos: &A) -> usize
157    where
158        A: GGLWEInfos,
159    {
160        assert_eq!(
161            infos.rank_in(),
162            infos.rank_out(),
163            "rank_in != rank_out is not supported for GGLWEAutomorphismKey"
164        );
165        GGLWESwitchingKey::alloc_bytes(infos)
166    }
167
168    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank: Rank, dnum: Dnum, dsize: Dsize) -> usize {
169        GGLWESwitchingKey::alloc_bytes_with(n, base2k, k, rank, rank, dnum, dsize)
170    }
171}
172
173impl<D: DataRef> GGLWEAutomorphismKey<D> {
174    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
175        self.key.at(row, col)
176    }
177}
178
179impl<D: DataMut> GGLWEAutomorphismKey<D> {
180    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
181        self.key.at_mut(row, col)
182    }
183}
184
185impl<D: DataMut> ReaderFrom for GGLWEAutomorphismKey<D> {
186    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
187        self.p = reader.read_u64::<LittleEndian>()? as i64;
188        self.key.read_from(reader)
189    }
190}
191
192impl<D: DataRef> WriterTo for GGLWEAutomorphismKey<D> {
193    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
194        writer.write_u64::<LittleEndian>(self.p as u64)?;
195        self.key.write_to(writer)
196    }
197}