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, 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 GLWESwitchingKeyCompressed<Vec<u8>> {
 85    pub fn alloc_from_infos<A>(infos: &A) -> Self
 86    where
 87        A: GGLWEInfos,
 88    {
 89        Self::alloc(
 90            infos.n(),
 91            infos.base2k(),
 92            infos.k(),
 93            infos.rank_in(),
 94            infos.rank_out(),
 95            infos.dnum(),
 96            infos.dsize(),
 97        )
 98    }
 99
100    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank_in: Rank, rank_out: Rank, dnum: Dnum, dsize: Dsize) -> Self {
101        GLWESwitchingKeyCompressed {
102            key: GGLWECompressed::alloc(n, base2k, k, rank_in, rank_out, dnum, dsize),
103            sk_in_n: 0,
104            sk_out_n: 0,
105        }
106    }
107
108    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
109    where
110        A: GGLWEInfos,
111    {
112        GGLWECompressed::bytes_of_from_infos(infos)
113    }
114
115    pub fn bytes_of(n: Degree, base2k: Base2K, k: TorusPrecision, rank_in: Rank, dnum: Dnum, dsize: Dsize) -> usize
116where {
117        GGLWECompressed::bytes_of(n, base2k, k, rank_in, dnum, dsize)
118    }
119}
120
121impl<D: DataMut> ReaderFrom for GLWESwitchingKeyCompressed<D> {
122    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
123        self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
124        self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
125        self.key.read_from(reader)
126    }
127}
128
129impl<D: DataRef> WriterTo for GLWESwitchingKeyCompressed<D> {
130    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
131        writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
132        writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
133        self.key.write_to(writer)
134    }
135}
136
137pub trait GLWESwitchingKeyDecompress
138where
139    Self: GGLWEDecompress,
140{
141    fn decompress_glwe_switching_key<R, O>(&self, res: &mut R, other: &O)
142    where
143        R: GLWESwitchingKeyToMut + GLWESwitchingKeySetMetaData,
144        O: GLWESwitchingKeyCompressedToRef,
145    {
146        let other: &GLWESwitchingKeyCompressed<&[u8]> = &other.to_ref();
147        self.decompress_gglwe(&mut res.to_mut().key, &other.key);
148        res.set_sk_in_n(other.sk_in_n);
149        res.set_sk_out_n(other.sk_out_n);
150    }
151}
152
153impl<B: Backend> GLWESwitchingKeyDecompress for Module<B> where Self: GGLWEDecompress {}
154
155impl<D: DataMut> GLWESwitchingKey<D> {
156    pub fn decompress<O, M>(&mut self, module: &M, other: &O)
157    where
158        O: GLWESwitchingKeyCompressedToRef,
159        M: GLWESwitchingKeyDecompress,
160    {
161        module.decompress_glwe_switching_key(self, other);
162    }
163}
164
165pub trait GLWESwitchingKeyCompressedToMut {
166    fn to_mut(&mut self) -> GLWESwitchingKeyCompressed<&mut [u8]>;
167}
168
169impl<D: DataMut> GLWESwitchingKeyCompressedToMut for GLWESwitchingKeyCompressed<D>
170where
171    GGLWECompressed<D>: GGLWECompressedToMut,
172{
173    fn to_mut(&mut self) -> GLWESwitchingKeyCompressed<&mut [u8]> {
174        GLWESwitchingKeyCompressed {
175            sk_in_n: self.sk_in_n,
176            sk_out_n: self.sk_out_n,
177            key: self.key.to_mut(),
178        }
179    }
180}
181
182pub trait GLWESwitchingKeyCompressedToRef {
183    fn to_ref(&self) -> GLWESwitchingKeyCompressed<&[u8]>;
184}
185
186impl<D: DataRef> GLWESwitchingKeyCompressedToRef for GLWESwitchingKeyCompressed<D>
187where
188    GGLWECompressed<D>: GGLWECompressedToRef,
189{
190    fn to_ref(&self) -> GLWESwitchingKeyCompressed<&[u8]> {
191        GLWESwitchingKeyCompressed {
192            sk_in_n: self.sk_in_n,
193            sk_out_n: self.sk_out_n,
194            key: self.key.to_ref(),
195        }
196    }
197}