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