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, B: Backend>(module: &Module<B>, infos: &A) -> Self
139    where
140        A: GGLWEInfos,
141        Module<B>: GLWESwitchingKeyCompressedAlloc,
142    {
143        module.alloc_glwe_switching_key_compressed_from_infos(infos)
144    }
145
146    pub fn alloc<B: Backend>(
147        module: &Module<B>,
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        Module<B>: 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, B: Backend>(module: &Module<B>, infos: &A) -> usize
162    where
163        A: GGLWEInfos,
164        Module<B>: GLWESwitchingKeyCompressedAlloc,
165    {
166        module.bytes_of_glwe_switching_key_compressed_from_infos(infos)
167    }
168
169    pub fn bytes_of<B: Backend>(
170        module: &Module<B>,
171        base2k: Base2K,
172        k: TorusPrecision,
173        rank_in: Rank,
174        dnum: Dnum,
175        dsize: Dsize,
176    ) -> usize
177    where
178        Module<B>: GLWESwitchingKeyCompressedAlloc,
179    {
180        module.bytes_of_glwe_switching_key_compressed(base2k, k, rank_in, dnum, dsize)
181    }
182}
183
184impl<D: DataMut> ReaderFrom for GLWESwitchingKeyCompressed<D> {
185    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
186        self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
187        self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
188        self.key.read_from(reader)
189    }
190}
191
192impl<D: DataRef> WriterTo for GLWESwitchingKeyCompressed<D> {
193    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
194        writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
195        writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
196        self.key.write_to(writer)
197    }
198}
199
200pub trait GLWESwitchingKeyDecompress
201where
202    Self: GGLWEDecompress,
203{
204    fn decompress_glwe_switching_key<R, O>(&self, res: &mut R, other: &O)
205    where
206        R: GLWESwitchingKeyToMut + GLWESwitchingKeySetMetaData,
207        O: GLWESwitchingKeyCompressedToRef,
208    {
209        let other: &GLWESwitchingKeyCompressed<&[u8]> = &other.to_ref();
210        self.decompress_gglwe(&mut res.to_mut().key, &other.key);
211        res.set_sk_in_n(other.sk_in_n);
212        res.set_sk_out_n(other.sk_out_n);
213    }
214}
215
216impl<B: Backend> GLWESwitchingKeyDecompress for Module<B> where Self: GGLWEDecompress {}
217
218impl<D: DataMut> GLWESwitchingKey<D> {
219    pub fn decompress<O, B: Backend>(&mut self, module: Module<B>, other: &O)
220    where
221        O: GLWESwitchingKeyCompressedToRef,
222        Module<B>: GGLWEDecompress,
223    {
224        module.decompress_glwe_switching_key(self, other);
225    }
226}
227
228pub trait GLWESwitchingKeyCompressedToMut {
229    fn to_mut(&mut self) -> GLWESwitchingKeyCompressed<&mut [u8]>;
230}
231
232impl<D: DataMut> GLWESwitchingKeyCompressedToMut for GLWESwitchingKeyCompressed<D>
233where
234    GGLWECompressed<D>: GGLWECompressedToMut,
235{
236    fn to_mut(&mut self) -> GLWESwitchingKeyCompressed<&mut [u8]> {
237        GLWESwitchingKeyCompressed {
238            sk_in_n: self.sk_in_n,
239            sk_out_n: self.sk_out_n,
240            key: self.key.to_mut(),
241        }
242    }
243}
244
245pub trait GLWESwitchingKeyCompressedToRef {
246    fn to_ref(&self) -> GLWESwitchingKeyCompressed<&[u8]>;
247}
248
249impl<D: DataRef> GLWESwitchingKeyCompressedToRef for GLWESwitchingKeyCompressed<D>
250where
251    GGLWECompressed<D>: GGLWECompressedToRef,
252{
253    fn to_ref(&self) -> GLWESwitchingKeyCompressed<&[u8]> {
254        GLWESwitchingKeyCompressed {
255            sk_in_n: self.sk_in_n,
256            sk_out_n: self.sk_out_n,
257            key: self.key.to_ref(),
258        }
259    }
260}