1use poulpy_hal::{
  2    layouts::{Data, DataMut, DataRef, FillUniform, ReaderFrom, WriterTo},
  3    source::Source,
  4};
  5
  6use crate::layouts::{
  7    Base2K, Degree, Dnum, Dsize, GGLWE, GGLWEInfos, GGLWEToMut, GGLWEToRef, GLWE, GLWEInfos, LWEInfos, Rank, TorusPrecision,
  8};
  9use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
 10
 11use std::fmt;
 12
 13#[derive(PartialEq, Eq, Copy, Clone, Debug)]
 14pub struct GLWESwitchingKeyLayout {
 15    pub n: Degree,
 16    pub base2k: Base2K,
 17    pub k: TorusPrecision,
 18    pub rank_in: Rank,
 19    pub rank_out: Rank,
 20    pub dnum: Dnum,
 21    pub dsize: Dsize,
 22}
 23
 24impl LWEInfos for GLWESwitchingKeyLayout {
 25    fn n(&self) -> Degree {
 26        self.n
 27    }
 28
 29    fn base2k(&self) -> Base2K {
 30        self.base2k
 31    }
 32
 33    fn k(&self) -> TorusPrecision {
 34        self.k
 35    }
 36}
 37
 38impl GLWEInfos for GLWESwitchingKeyLayout {
 39    fn rank(&self) -> Rank {
 40        self.rank_out()
 41    }
 42}
 43
 44impl GGLWEInfos for GLWESwitchingKeyLayout {
 45    fn rank_in(&self) -> Rank {
 46        self.rank_in
 47    }
 48
 49    fn rank_out(&self) -> Rank {
 50        self.rank_out
 51    }
 52
 53    fn dsize(&self) -> Dsize {
 54        self.dsize
 55    }
 56
 57    fn dnum(&self) -> Dnum {
 58        self.dnum
 59    }
 60}
 61
 62#[derive(PartialEq, Eq, Clone)]
 63pub struct GLWESwitchingKey<D: Data> {
 64    pub(crate) key: GGLWE<D>,
 65    pub(crate) input_degree: Degree,  // Degree of sk_in
 66    pub(crate) output_degree: Degree, // Degree of sk_out
 67}
 68
 69pub trait GLWESwitchingKeyDegrees {
 70    fn input_degree(&self) -> &Degree;
 71    fn output_degree(&self) -> &Degree;
 72}
 73
 74impl<D: DataRef> GLWESwitchingKeyDegrees for GLWESwitchingKey<D> {
 75    fn output_degree(&self) -> &Degree {
 76        &self.output_degree
 77    }
 78
 79    fn input_degree(&self) -> &Degree {
 80        &self.input_degree
 81    }
 82}
 83
 84pub trait GLWESwitchingKeyDegreesMut {
 85    fn input_degree(&mut self) -> &mut Degree;
 86    fn output_degree(&mut self) -> &mut Degree;
 87}
 88
 89impl<D: DataMut> GLWESwitchingKeyDegreesMut for GLWESwitchingKey<D> {
 90    fn output_degree(&mut self) -> &mut Degree {
 91        &mut self.output_degree
 92    }
 93
 94    fn input_degree(&mut self) -> &mut Degree {
 95        &mut self.input_degree
 96    }
 97}
 98
 99impl<D: Data> LWEInfos for GLWESwitchingKey<D> {
100    fn n(&self) -> Degree {
101        self.key.n()
102    }
103
104    fn base2k(&self) -> Base2K {
105        self.key.base2k()
106    }
107
108    fn k(&self) -> TorusPrecision {
109        self.key.k()
110    }
111
112    fn size(&self) -> usize {
113        self.key.size()
114    }
115}
116
117impl<D: Data> GLWEInfos for GLWESwitchingKey<D> {
118    fn rank(&self) -> Rank {
119        self.rank_out()
120    }
121}
122
123impl<D: Data> GGLWEInfos for GLWESwitchingKey<D> {
124    fn rank_in(&self) -> Rank {
125        self.key.rank_in()
126    }
127
128    fn rank_out(&self) -> Rank {
129        self.key.rank_out()
130    }
131
132    fn dsize(&self) -> Dsize {
133        self.key.dsize()
134    }
135
136    fn dnum(&self) -> Dnum {
137        self.key.dnum()
138    }
139}
140
141impl<D: DataRef> fmt::Debug for GLWESwitchingKey<D> {
142    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
143        write!(f, "{self}")
144    }
145}
146
147impl<D: DataRef> fmt::Display for GLWESwitchingKey<D> {
148    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
149        write!(
150            f,
151            "(GLWESwitchingKey: sk_in_n={} sk_out_n={}) {}",
152            self.input_degree,
153            self.output_degree,
154            self.key.data()
155        )
156    }
157}
158
159impl<D: DataMut> FillUniform for GLWESwitchingKey<D> {
160    fn fill_uniform(&mut self, log_bound: usize, source: &mut Source) {
161        self.key.fill_uniform(log_bound, source);
162    }
163}
164
165impl GLWESwitchingKey<Vec<u8>> {
166    pub fn alloc_from_infos<A>(infos: &A) -> Self
167    where
168        A: GGLWEInfos,
169    {
170        Self::alloc(
171            infos.n(),
172            infos.base2k(),
173            infos.k(),
174            infos.rank_in(),
175            infos.rank_out(),
176            infos.dnum(),
177            infos.dsize(),
178        )
179    }
180
181    pub fn alloc(n: Degree, base2k: Base2K, k: TorusPrecision, rank_in: Rank, rank_out: Rank, dnum: Dnum, dsize: Dsize) -> Self {
182        GLWESwitchingKey {
183            key: GGLWE::alloc(n, base2k, k, rank_in, rank_out, dnum, dsize),
184            input_degree: Degree(0),
185            output_degree: Degree(0),
186        }
187    }
188
189    pub fn bytes_of_from_infos<A>(infos: &A) -> usize
190    where
191        A: GGLWEInfos,
192    {
193        Self::bytes_of(
194            infos.n(),
195            infos.base2k(),
196            infos.k(),
197            infos.rank_in(),
198            infos.rank_out(),
199            infos.dnum(),
200            infos.dsize(),
201        )
202    }
203
204    pub fn bytes_of(
205        n: Degree,
206        base2k: Base2K,
207        k: TorusPrecision,
208        rank_in: Rank,
209        rank_out: Rank,
210        dnum: Dnum,
211        dsize: Dsize,
212    ) -> usize {
213        GGLWE::bytes_of(n, base2k, k, rank_in, rank_out, dnum, dsize)
214    }
215}
216
217impl<D: DataMut> GGLWEToMut for GLWESwitchingKey<D> {
218    fn to_mut(&mut self) -> GGLWE<&mut [u8]> {
219        self.key.to_mut()
220    }
221}
222
223impl<D: DataRef> GGLWEToRef for GLWESwitchingKey<D> {
224    fn to_ref(&self) -> GGLWE<&[u8]> {
225        self.key.to_ref()
226    }
227}
228
229impl<D: DataRef> GLWESwitchingKey<D> {
230    pub fn at(&self, row: usize, col: usize) -> GLWE<&[u8]> {
231        self.key.at(row, col)
232    }
233}
234
235impl<D: DataMut> GLWESwitchingKey<D> {
236    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWE<&mut [u8]> {
237        self.key.at_mut(row, col)
238    }
239}
240
241impl<D: DataMut> ReaderFrom for GLWESwitchingKey<D> {
242    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
243        self.input_degree = Degree(reader.read_u32::<LittleEndian>()? as u32);
244        self.output_degree = Degree(reader.read_u32::<LittleEndian>()? as u32);
245        self.key.read_from(reader)
246    }
247}
248
249impl<D: DataRef> WriterTo for GLWESwitchingKey<D> {
250    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
251        writer.write_u32::<LittleEndian>(self.input_degree.into())?;
252        writer.write_u32::<LittleEndian>(self.output_degree.into())?;
253        self.key.write_to(writer)
254    }
255}