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