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