1use poulpy_hal::{
  2    layouts::{Backend, Data, DataMut, DataRef, FillUniform, Module, ReaderFrom, WriterTo},
  3    source::Source,
  4};
  5
  6use crate::layouts::{
  7    Base2K, Dnum, Dsize, GGLWE, GGLWEAlloc, GGLWEInfos, GGLWEToMut, GGLWEToRef, GLWE, GLWEInfos, LWEInfos, Rank, RingDegree,
  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: RingDegree,
 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) -> RingDegree {
 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 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 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) -> RingDegree {
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<B: Backend> GLWESwitchingKeyAlloc for Module<B> where Self: GGLWEAlloc {}
228
229impl GLWESwitchingKey<Vec<u8>> {
230    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
231    where
232        A: GGLWEInfos,
233        M: GLWESwitchingKeyAlloc,
234    {
235        module.alloc_glwe_switching_key_from_infos(infos)
236    }
237
238    pub fn alloc<M>(
239        module: &M,
240        base2k: Base2K,
241        k: TorusPrecision,
242        rank_in: Rank,
243        rank_out: Rank,
244        dnum: Dnum,
245        dsize: Dsize,
246    ) -> Self
247    where
248        M: GLWESwitchingKeyAlloc,
249    {
250        module.alloc_glwe_switching_key(base2k, k, rank_in, rank_out, dnum, dsize)
251    }
252
253    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
254    where
255        A: GGLWEInfos,
256        M: GLWESwitchingKeyAlloc,
257    {
258        module.bytes_of_glwe_switching_key_from_infos(infos)
259    }
260
261    pub fn bytes_of<M>(
262        module: &M,
263        base2k: Base2K,
264        k: TorusPrecision,
265        rank_in: Rank,
266        rank_out: Rank,
267        dnum: Dnum,
268        dsize: Dsize,
269    ) -> usize
270    where
271        M: GLWESwitchingKeyAlloc,
272    {
273        module.bytes_of_glwe_switching_key(base2k, k, rank_in, rank_out, dnum, dsize)
274    }
275}
276
277pub trait GLWESwitchingKeyToMut {
278    fn to_mut(&mut self) -> GLWESwitchingKey<&mut [u8]>;
279}
280
281impl<D: DataMut> GLWESwitchingKeyToMut for GLWESwitchingKey<D>
282where
283    GGLWE<D>: GGLWEToMut,
284{
285    fn to_mut(&mut self) -> GLWESwitchingKey<&mut [u8]> {
286        GLWESwitchingKey {
287            key: self.key.to_mut(),
288            sk_in_n: self.sk_in_n,
289            sk_out_n: self.sk_out_n,
290        }
291    }
292}
293
294pub trait GLWESwitchingKeyToRef {
295    fn to_ref(&self) -> GLWESwitchingKey<&[u8]>;
296}
297
298impl<D: DataRef> GLWESwitchingKeyToRef for GLWESwitchingKey<D>
299where
300    GGLWE<D>: GGLWEToRef,
301{
302    fn to_ref(&self) -> GLWESwitchingKey<&[u8]> {
303        GLWESwitchingKey {
304            key: self.key.to_ref(),
305            sk_in_n: self.sk_in_n,
306            sk_out_n: self.sk_out_n,
307        }
308    }
309}
310
311impl<D: DataRef> GLWESwitchingKey<D> {
312    pub fn at(&self, row: usize, col: usize) -> GLWE<&[u8]> {
313        self.key.at(row, col)
314    }
315}
316
317impl<D: DataMut> GLWESwitchingKey<D> {
318    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWE<&mut [u8]> {
319        self.key.at_mut(row, col)
320    }
321}
322
323impl<D: DataMut> ReaderFrom for GLWESwitchingKey<D> {
324    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
325        self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
326        self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
327        self.key.read_from(reader)
328    }
329}
330
331impl<D: DataRef> WriterTo for GLWESwitchingKey<D> {
332    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
333        writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
334        writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
335        self.key.write_to(writer)
336    }
337}