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};
  9use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 10
 11use std::fmt;
 12
 13#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 14pub struct GLWESwitchingKeyLayout {
 15    pub n: Degree,
 16    pub base2k: Base2K,
 17    pub k: TorusPrecision,
 18    pub rank_in: Rank,
 19    pub rank_out: Rank,
 20    pub dnum: Dnum,
 21    pub dsize: Dsize,
 22}
 23
 24impl LWEInfos for GLWESwitchingKeyLayout {
 25    fn n(&self) -> Degree {
 26        self.n
 27    }
 28
 29    fn base2k(&self) -> Base2K {
 30        self.base2k
 31    }
 32
 33    fn k(&self) -> TorusPrecision {
 34        self.k
 35    }
 36}
 37
 38impl GLWEInfos for GLWESwitchingKeyLayout {
 39    fn rank(&self) -> Rank {
 40        self.rank_out()
 41    }
 42}
 43
 44impl GGLWEInfos for GLWESwitchingKeyLayout {
 45    fn rank_in(&self) -> Rank {
 46        self.rank_in
 47    }
 48
 49    fn rank_out(&self) -> Rank {
 50        self.rank_out
 51    }
 52
 53    fn dsize(&self) -> Dsize {
 54        self.dsize
 55    }
 56
 57    fn dnum(&self) -> Dnum {
 58        self.dnum
 59    }
 60}
 61
 62#[derive(PartialEq, Eq, Clone)]
 63pub struct GLWESwitchingKey<D: Data> {
 64    pub(crate) key: GGLWE<D>,
 65    pub(crate) sk_in_n: usize,  // Degree of sk_in
 66    pub(crate) sk_out_n: usize, // Degree of sk_out
 67}
 68
 69pub trait GLWESwitchingKeySetMetaData {
 70    fn set_sk_in_n(&mut self, sk_in_n: usize);
 71    fn set_sk_out_n(&mut self, sk_out_n: usize);
 72}
 73
 74impl<D: DataMut> GLWESwitchingKeySetMetaData for GLWESwitchingKey<D> {
 75    fn set_sk_in_n(&mut self, sk_in_n: usize) {
 76        self.sk_in_n = sk_in_n
 77    }
 78
 79    fn set_sk_out_n(&mut self, sk_out_n: usize) {
 80        self.sk_out_n = sk_out_n
 81    }
 82}
 83
 84pub trait GLWESwtichingKeyGetMetaData {
 85    fn sk_in_n(&self) -> usize;
 86    fn sk_out_n(&self) -> usize;
 87}
 88
 89impl<D: DataRef> GLWESwtichingKeyGetMetaData for GLWESwitchingKey<D> {
 90    fn sk_in_n(&self) -> usize {
 91        self.sk_in_n
 92    }
 93
 94    fn sk_out_n(&self) -> usize {
 95        self.sk_out_n
 96    }
 97}
 98
 99impl<D: Data> LWEInfos for GLWESwitchingKey<D> {
100    fn n(&self) -> Degree {
101        self.key.n()
102    }
103
104    fn base2k(&self) -> Base2K {
105        self.key.base2k()
106    }
107
108    fn k(&self) -> TorusPrecision {
109        self.key.k()
110    }
111
112    fn size(&self) -> usize {
113        self.key.size()
114    }
115}
116
117impl<D: Data> GLWEInfos for GLWESwitchingKey<D> {
118    fn rank(&self) -> Rank {
119        self.rank_out()
120    }
121}
122
123impl<D: Data> GGLWEInfos for GLWESwitchingKey<D> {
124    fn rank_in(&self) -> Rank {
125        self.key.rank_in()
126    }
127
128    fn rank_out(&self) -> Rank {
129        self.key.rank_out()
130    }
131
132    fn dsize(&self) -> Dsize {
133        self.key.dsize()
134    }
135
136    fn dnum(&self) -> Dnum {
137        self.key.dnum()
138    }
139}
140
141impl<D: DataRef> fmt::Debug for GLWESwitchingKey<D> {
142    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
143        write!(f, "{self}")
144    }
145}
146
147impl<D: DataRef> fmt::Display for GLWESwitchingKey<D> {
148    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
149        write!(
150            f,
151            "(GLWESwitchingKey: sk_in_n={} sk_out_n={}) {}",
152            self.sk_in_n,
153            self.sk_out_n,
154            self.key.data()
155        )
156    }
157}
158
159impl<D: DataMut> FillUniform for GLWESwitchingKey<D> {
160    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
161        self.key.fill_uniform(log_bound, source);
162    }
163}
164
165impl GLWESwitchingKey<Vec<u8>> {
166    pub fn alloc_from_infos<A>(infos: &A) -> Self
167    where
168        A: GGLWEInfos,
169    {
170        Self::alloc(
171            infos.n(),
172            infos.base2k(),
173            infos.k(),
174            infos.rank_in(),
175            infos.rank_out(),
176            infos.dnum(),
177            infos.dsize(),
178        )
179    }
180
181    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank_in: Rank, rank_out: Rank, dnum: Dnum, dsize: Dsize) -> Self {
182        GLWESwitchingKey {
183            key: GGLWE::alloc(n, base2k, k, rank_in, rank_out, dnum, dsize),
184            sk_in_n: 0,
185            sk_out_n: 0,
186        }
187    }
188
189    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
190    where
191        A: GGLWEInfos,
192    {
193        Self::bytes_of(
194            infos.n(),
195            infos.base2k(),
196            infos.k(),
197            infos.rank_in(),
198            infos.rank_out(),
199            infos.dnum(),
200            infos.dsize(),
201        )
202    }
203
204    pub fn bytes_of(
205        n: Degree,
206        base2k: Base2K,
207        k: TorusPrecision,
208        rank_in: Rank,
209        rank_out: Rank,
210        dnum: Dnum,
211        dsize: Dsize,
212    ) -> usize {
213        GGLWE::bytes_of(n, base2k, k, rank_in, rank_out, dnum, dsize)
214    }
215}
216
217pub trait GLWESwitchingKeyToMut {
218    fn to_mut(&mut self) -> GLWESwitchingKey<&mut [u8]>;
219}
220
221impl<D: DataMut> GLWESwitchingKeyToMut for GLWESwitchingKey<D>
222where
223    GGLWE<D>: GGLWEToMut,
224{
225    fn to_mut(&mut self) -> GLWESwitchingKey<&mut [u8]> {
226        GLWESwitchingKey {
227            key: self.key.to_mut(),
228            sk_in_n: self.sk_in_n,
229            sk_out_n: self.sk_out_n,
230        }
231    }
232}
233
234pub trait GLWESwitchingKeyToRef {
235    fn to_ref(&self) -> GLWESwitchingKey<&[u8]>;
236}
237
238impl<D: DataRef> GLWESwitchingKeyToRef for GLWESwitchingKey<D>
239where
240    GGLWE<D>: GGLWEToRef,
241{
242    fn to_ref(&self) -> GLWESwitchingKey<&[u8]> {
243        GLWESwitchingKey {
244            key: self.key.to_ref(),
245            sk_in_n: self.sk_in_n,
246            sk_out_n: self.sk_out_n,
247        }
248    }
249}
250
251impl<D: DataRef> GLWESwitchingKey<D> {
252    pub fn at(&self, row: usize, col: usize) -> GLWE<&[u8]> {
253        self.key.at(row, col)
254    }
255}
256
257impl<D: DataMut> GLWESwitchingKey<D> {
258    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWE<&mut [u8]> {
259        self.key.at_mut(row, col)
260    }
261}
262
263impl<D: DataMut> ReaderFrom for GLWESwitchingKey<D> {
264    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
265        self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
266        self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
267        self.key.read_from(reader)
268    }
269}
270
271impl<D: DataRef> WriterTo for GLWESwitchingKey<D> {
272    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
273        writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
274        writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
275        self.key.write_to(writer)
276    }
277}