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, GLWESwitchingKeyCompressedToMut, GLWESwitchingKeyCompressedToRef, GLWESwitchingKeyDecompress,
 11    },
 12};
 13use std::fmt;
 14
 15#[derive(PartialEq, Eq, Clone)]
 16pub struct LWEToGLWESwitchingKeyCompressed<D: Data>(pub(crate) GLWESwitchingKeyCompressed<D>);
 17
 18impl<D: Data> LWEInfos for LWEToGLWESwitchingKeyCompressed<D> {
 19    fn n(&self) -> Degree {
 20        self.0.n()
 21    }
 22
 23    fn base2k(&self) -> Base2K {
 24        self.0.base2k()
 25    }
 26
 27    fn k(&self) -> TorusPrecision {
 28        self.0.k()
 29    }
 30    fn size(&self) -> usize {
 31        self.0.size()
 32    }
 33}
 34impl<D: Data> GLWEInfos for LWEToGLWESwitchingKeyCompressed<D> {
 35    fn rank(&self) -> Rank {
 36        self.rank_out()
 37    }
 38}
 39
 40impl<D: Data> GGLWEInfos for LWEToGLWESwitchingKeyCompressed<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 LWEToGLWESwitchingKeyCompressed<D> {
 59    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
 60        write!(f, "{self}")
 61    }
 62}
 63
 64impl<D: DataMut> FillUniform for LWEToGLWESwitchingKeyCompressed<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 LWEToGLWESwitchingKeyCompressed<D> {
 71    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 72        write!(f, "(LWEToGLWESwitchingKeyCompressed) {}", self.0)
 73    }
 74}
 75
 76impl<D: DataMut> ReaderFrom for LWEToGLWESwitchingKeyCompressed<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 LWEToGLWESwitchingKeyCompressed<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
 88impl LWEToGLWESwitchingKeyCompressed<Vec<u8>> {
 89    pub fn alloc_from_infos<A>(infos: &A) -> Self
 90    where
 91        A: GGLWEInfos,
 92    {
 93        assert_eq!(
 94            infos.dsize().0,
 95            1,
 96            "dsize > 1 is not supported for LWEToGLWESwitchingKeyCompressed"
 97        );
 98        assert_eq!(
 99            infos.rank_in().0,
100            1,
101            "rank_in > 1 is not supported for LWEToGLWESwitchingKeyCompressed"
102        );
103        Self::alloc(
104            infos.n(),
105            infos.base2k(),
106            infos.k(),
107            infos.rank_out(),
108            infos.dnum(),
109        )
110    }
111
112    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank_out: Rank, dnum: Dnum) -> Self {
113        LWEToGLWESwitchingKeyCompressed(GLWESwitchingKeyCompressed::alloc(
114            n,
115            base2k,
116            k,
117            Rank(1),
118            rank_out,
119            dnum,
120            Dsize(1),
121        ))
122    }
123
124    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
125    where
126        A: GGLWEInfos,
127    {
128        assert_eq!(
129            infos.dsize().0,
130            1,
131            "dsize > 1 is not supported for LWEToGLWESwitchingKeyCompressed"
132        );
133        assert_eq!(
134            infos.rank_in().0,
135            1,
136            "rank_in > 1 is not supported for LWEToGLWESwitchingKeyCompressed"
137        );
138        GLWESwitchingKeyCompressed::bytes_of_from_infos(infos)
139    }
140
141    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, dnum: Dnum) -> usize {
142        GLWESwitchingKeyCompressed::bytes_of(n, base2k, k, Rank(1), dnum, Dsize(1))
143    }
144}
145
146pub trait LWEToGLWESwitchingKeyDecompress
147where
148    Self: GLWESwitchingKeyDecompress,
149{
150    fn decompress_lwe_to_glwe_switching_key<R, O>(&self, res: &mut R, other: &O)
151    where
152        R: LWEToGLWESwitchingKeyToMut,
153        O: LWEToGLWESwitchingKeyCompressedToRef,
154    {
155        self.decompress_glwe_switching_key(&mut res.to_mut().0, &other.to_ref().0);
156    }
157}
158
159impl<B: Backend> LWEToGLWESwitchingKeyDecompress for Module<B> where Self: GLWESwitchingKeyDecompress {}
160
161impl<D: DataMut> LWEToGLWESwitchingKey<D> {
162    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
163    where
164        O: LWEToGLWESwitchingKeyCompressedToRef,
165        M: LWEToGLWESwitchingKeyDecompress,
166    {
167        module.decompress_lwe_to_glwe_switching_key(self, other);
168    }
169}
170
171pub trait LWEToGLWESwitchingKeyCompressedToRef {
172    fn to_ref(&self) -> LWEToGLWESwitchingKeyCompressed<&[u8]>;
173}
174
175impl<D: DataRef> LWEToGLWESwitchingKeyCompressedToRef for LWEToGLWESwitchingKeyCompressed<D>
176where
177    GLWESwitchingKeyCompressed<D>: GLWESwitchingKeyCompressedToRef,
178{
179    fn to_ref(&self) -> LWEToGLWESwitchingKeyCompressed<&[u8]> {
180        LWEToGLWESwitchingKeyCompressed(self.0.to_ref())
181    }
182}
183
184pub trait LWEToGLWESwitchingKeyCompressedToMut {
185    fn to_mut(&mut self) -> LWEToGLWESwitchingKeyCompressed<&mut [u8]>;
186}
187
188impl<D: DataMut> LWEToGLWESwitchingKeyCompressedToMut for LWEToGLWESwitchingKeyCompressed<D>
189where
190    GLWESwitchingKeyCompressed<D>: GLWESwitchingKeyCompressedToMut,
191{
192    fn to_mut(&mut self) -> LWEToGLWESwitchingKeyCompressed<&mut [u8]> {
193        LWEToGLWESwitchingKeyCompressed(self.0.to_mut())
194    }
195}