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