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