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
 84impl<BE: Backend> GLWESwitchingKeyCompressedAlloc for Module<BE> where Self: GGLWECompressedAlloc {}
 85
 86pub trait GLWESwitchingKeyCompressedAlloc
 87where
 88    Self: GGLWECompressedAlloc,
 89{
 90    fn alloc_glwe_switching_key_compressed(
 91        &self,
 92        base2k: Base2K,
 93        k: TorusPrecision,
 94        rank_in: Rank,
 95        rank_out: Rank,
 96        dnum: Dnum,
 97        dsize: Dsize,
 98    ) -> GLWESwitchingKeyCompressed<Vec<u8>> {
 99        GLWESwitchingKeyCompressed {
100            key: self.alloc_gglwe_compressed(base2k, k, rank_in, rank_out, dnum, dsize),
101            sk_in_n: 0,
102            sk_out_n: 0,
103        }
104    }
105
106    fn alloc_glwe_switching_key_compressed_from_infos<A>(&self, infos: &A) -> GLWESwitchingKeyCompressed<Vec<u8>>
107    where
108        A: GGLWEInfos,
109    {
110        self.alloc_glwe_switching_key_compressed(
111            infos.base2k(),
112            infos.k(),
113            infos.rank_in(),
114            infos.rank_out(),
115            infos.dnum(),
116            infos.dsize(),
117        )
118    }
119
120    fn bytes_of_glwe_switching_key_compressed(
121        &self,
122        base2k: Base2K,
123        k: TorusPrecision,
124        rank_in: Rank,
125        dnum: Dnum,
126        dsize: Dsize,
127    ) -> usize {
128        self.bytes_of_gglwe_compressed(base2k, k, rank_in, dnum, dsize)
129    }
130
131    fn bytes_of_glwe_switching_key_compressed_from_infos<A>(&self, infos: &A) -> usize
132    where
133        A: GGLWEInfos,
134    {
135        self.bytes_of_gglwe_compressed_from_infos(infos)
136    }
137}
138
139impl GLWESwitchingKeyCompressed<Vec<u8>> {
140    pub fn alloc_from_infos<A, M>(module: &M, infos: &A) -> Self
141    where
142        A: GGLWEInfos,
143        M: GLWESwitchingKeyCompressedAlloc,
144    {
145        module.alloc_glwe_switching_key_compressed_from_infos(infos)
146    }
147
148    pub fn alloc<M>(
149        module: &M,
150        base2k: Base2K,
151        k: TorusPrecision,
152        rank_in: Rank,
153        rank_out: Rank,
154        dnum: Dnum,
155        dsize: Dsize,
156    ) -> Self
157    where
158        M: GLWESwitchingKeyCompressedAlloc,
159    {
160        module.alloc_glwe_switching_key_compressed(base2k, k, rank_in, rank_out, dnum, dsize)
161    }
162
163    pub fn bytes_of_from_infos<A, M>(module: &M, infos: &A) -> usize
164    where
165        A: GGLWEInfos,
166        M: GLWESwitchingKeyCompressedAlloc,
167    {
168        module.bytes_of_glwe_switching_key_compressed_from_infos(infos)
169    }
170
171    pub fn bytes_of<M>(module: &M, base2k: Base2K, k: TorusPrecision, rank_in: Rank, dnum: Dnum, dsize: Dsize) -> usize
172    where
173        M: GLWESwitchingKeyCompressedAlloc,
174    {
175        module.bytes_of_glwe_switching_key_compressed(base2k, k, rank_in, dnum, dsize)
176    }
177}
178
179impl<D: DataMut> ReaderFrom for GLWESwitchingKeyCompressed<D> {
180    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
181        self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
182        self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
183        self.key.read_from(reader)
184    }
185}
186
187impl<D: DataRef> WriterTo for GLWESwitchingKeyCompressed<D> {
188    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
189        writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
190        writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
191        self.key.write_to(writer)
192    }
193}
194
195pub trait GLWESwitchingKeyDecompress
196where
197    Self: GGLWEDecompress,
198{
199    fn decompress_glwe_switching_key<R, O>(&self, res: &mut R, other: &O)
200    where
201        R: GLWESwitchingKeyToMut + GLWESwitchingKeySetMetaData,
202        O: GLWESwitchingKeyCompressedToRef,
203    {
204        let other: &GLWESwitchingKeyCompressed<&[u8]> = &other.to_ref();
205        self.decompress_gglwe(&mut res.to_mut().key, &other.key);
206        res.set_sk_in_n(other.sk_in_n);
207        res.set_sk_out_n(other.sk_out_n);
208    }
209}
210
211impl<B: Backend> GLWESwitchingKeyDecompress for Module<B> where Self: GGLWEDecompress {}
212
213impl<D: DataMut> GLWESwitchingKey<D> {
214    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
215    where
216        O: GLWESwitchingKeyCompressedToRef,
217        M: GLWESwitchingKeyDecompress,
218    {
219        module.decompress_glwe_switching_key(self, other);
220    }
221}
222
223pub trait GLWESwitchingKeyCompressedToMut {
224    fn to_mut(&mut self) -> GLWESwitchingKeyCompressed<&mut [u8]>;
225}
226
227impl<D: DataMut> GLWESwitchingKeyCompressedToMut for GLWESwitchingKeyCompressed<D>
228where
229    GGLWECompressed<D>: GGLWECompressedToMut,
230{
231    fn to_mut(&mut self) -> GLWESwitchingKeyCompressed<&mut [u8]> {
232        GLWESwitchingKeyCompressed {
233            sk_in_n: self.sk_in_n,
234            sk_out_n: self.sk_out_n,
235            key: self.key.to_mut(),
236        }
237    }
238}
239
240pub trait GLWESwitchingKeyCompressedToRef {
241    fn to_ref(&self) -> GLWESwitchingKeyCompressed<&[u8]>;
242}
243
244impl<D: DataRef> GLWESwitchingKeyCompressedToRef for GLWESwitchingKeyCompressed<D>
245where
246    GGLWECompressed<D>: GGLWECompressedToRef,
247{
248    fn to_ref(&self) -> GLWESwitchingKeyCompressed<&[u8]> {
249        GLWESwitchingKeyCompressed {
250            sk_in_n: self.sk_in_n,
251            sk_out_n: self.sk_out_n,
252            key: self.key.to_ref(),
253        }
254    }
255}