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, GGLWECiphertextToMut, GGLWEInfos, GGLWEToRef, GLWECiphertext, 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
166impl GLWESwitchingKey<Vec<u8>> {
167    pub fn alloc<A>(infos: &A) -> Self
168    where
169        A: GGLWEInfos,
170    {
171        GLWESwitchingKey {
172            key: GGLWE::alloc(infos),
173            sk_in_n: 0,
174            sk_out_n: 0,
175        }
176    }
177
178    pub fn alloc_with(
179        n: Degree,
180        base2k: Base2K,
181        k: TorusPrecision,
182        rank_in: Rank,
183        rank_out: Rank,
184        dnum: Dnum,
185        dsize: Dsize,
186    ) -> Self {
187        GLWESwitchingKey {
188            key: GGLWE::alloc_with(n, base2k, k, rank_in, rank_out, dnum, dsize),
189            sk_in_n: 0,
190            sk_out_n: 0,
191        }
192    }
193
194    pub fn alloc_bytes<A>(infos: &A) -> usize
195    where
196        A: GGLWEInfos,
197    {
198        GGLWE::alloc_bytes(infos)
199    }
200
201    pub fn alloc_bytes_with(
202        n: Degree,
203        base2k: Base2K,
204        k: TorusPrecision,
205        rank_in: Rank,
206        rank_out: Rank,
207        dnum: Dnum,
208        dsize: Dsize,
209    ) -> usize {
210        GGLWE::alloc_bytes_with(n, base2k, k, rank_in, rank_out, dnum, dsize)
211    }
212}
213
214pub trait GLWESwitchingKeyToMut {
215    fn to_mut(&mut self) -> GLWESwitchingKey<&mut [u8]>;
216}
217
218impl<D: DataMut> GLWESwitchingKeyToMut for GLWESwitchingKey<D>
219where
220    GGLWE<D>: GGLWECiphertextToMut,
221{
222    fn to_mut(&mut self) -> GLWESwitchingKey<&mut [u8]> {
223        GLWESwitchingKey {
224            key: self.key.to_mut(),
225            sk_in_n: self.sk_in_n,
226            sk_out_n: self.sk_out_n,
227        }
228    }
229}
230
231pub trait GLWESwitchingKeyToRef {
232    fn to_ref(&self) -> GLWESwitchingKey<&[u8]>;
233}
234
235impl<D: DataRef> GLWESwitchingKeyToRef for GLWESwitchingKey<D>
236where
237    GGLWE<D>: GGLWEToRef,
238{
239    fn to_ref(&self) -> GLWESwitchingKey<&[u8]> {
240        GLWESwitchingKey {
241            key: self.key.to_ref(),
242            sk_in_n: self.sk_in_n,
243            sk_out_n: self.sk_out_n,
244        }
245    }
246}
247
248impl<D: DataRef> GLWESwitchingKey<D> {
249    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
250        self.key.at(row, col)
251    }
252}
253
254impl<D: DataMut> GLWESwitchingKey<D> {
255    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
256        self.key.at_mut(row, col)
257    }
258}
259
260impl<D: DataMut> ReaderFrom for GLWESwitchingKey<D> {
261    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
262        self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
263        self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
264        self.key.read_from(reader)
265    }
266}
267
268impl<D: DataRef> WriterTo for GLWESwitchingKey<D> {
269    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
270        writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
271        writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
272        self.key.write_to(writer)
273    }
274}