1use backend::hal::{
  2    api::{MatZnxAlloc, MatZnxAllocBytes},
  3    layouts::{Backend, Data, DataMut, DataRef, MatZnx, Module, ReaderFrom, Scratch, VmpPMat, WriterTo},
  4};
  5
  6use crate::{GGLWECiphertext, GGLWECiphertextExec, GGLWEExecLayoutFamily, GLWECiphertext, Infos};
  7
  8#[derive(PartialEq, Eq)]
  9pub struct GLWESwitchingKey<D: Data> {
 10    pub(crate) key: GGLWECiphertext<D>,
 11    pub(crate) sk_in_n: usize,  // Degree of sk_in
 12    pub(crate) sk_out_n: usize, // Degree of sk_out
 13}
 14
 15impl GLWESwitchingKey<Vec<u8>> {
 16    pub fn alloc<B: Backend>(
 17        module: &Module<B>,
 18        basek: usize,
 19        k: usize,
 20        rows: usize,
 21        digits: usize,
 22        rank_in: usize,
 23        rank_out: usize,
 24    ) -> Self
 25    where
 26        Module<B>: MatZnxAlloc,
 27    {
 28        GLWESwitchingKey {
 29            key: GGLWECiphertext::alloc(module, basek, k, rows, digits, rank_in, rank_out),
 30            sk_in_n: 0,
 31            sk_out_n: 0,
 32        }
 33    }
 34
 35    pub fn bytes_of<B: Backend>(
 36        module: &Module<B>,
 37        basek: usize,
 38        k: usize,
 39        rows: usize,
 40        digits: usize,
 41        rank_in: usize,
 42        rank_out: usize,
 43    ) -> usize
 44    where
 45        Module<B>: MatZnxAllocBytes,
 46    {
 47        GGLWECiphertext::<Vec<u8>>::bytes_of(module, basek, k, rows, digits, rank_in, rank_out)
 48    }
 49}
 50
 51impl<D: Data> Infos for GLWESwitchingKey<D> {
 52    type Inner = MatZnx<D>;
 53
 54    fn inner(&self) -> &Self::Inner {
 55        self.key.inner()
 56    }
 57
 58    fn basek(&self) -> usize {
 59        self.key.basek()
 60    }
 61
 62    fn k(&self) -> usize {
 63        self.key.k()
 64    }
 65}
 66
 67impl<D: Data> GLWESwitchingKey<D> {
 68    pub fn rank(&self) -> usize {
 69        self.key.data.cols_out() - 1
 70    }
 71
 72    pub fn rank_in(&self) -> usize {
 73        self.key.data.cols_in()
 74    }
 75
 76    pub fn rank_out(&self) -> usize {
 77        self.key.data.cols_out() - 1
 78    }
 79
 80    pub fn digits(&self) -> usize {
 81        self.key.digits()
 82    }
 83
 84    pub fn sk_degree_in(&self) -> usize {
 85        self.sk_in_n
 86    }
 87
 88    pub fn sk_degree_out(&self) -> usize {
 89        self.sk_out_n
 90    }
 91}
 92
 93impl<D: DataRef> GLWESwitchingKey<D> {
 94    pub fn at(&self, row: usize, col: usize) -> GLWECiphertext<&[u8]> {
 95        self.key.at(row, col)
 96    }
 97}
 98
 99impl<D: DataMut> GLWESwitchingKey<D> {
100    pub fn at_mut(&mut self, row: usize, col: usize) -> GLWECiphertext<&mut [u8]> {
101        self.key.at_mut(row, col)
102    }
103}
104
105use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
106
107impl<D: DataMut> ReaderFrom for GLWESwitchingKey<D> {
108    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
109        self.sk_in_n = reader.read_u64::<LittleEndian>()? as usize;
110        self.sk_out_n = reader.read_u64::<LittleEndian>()? as usize;
111        self.key.read_from(reader)
112    }
113}
114
115impl<D: DataRef> WriterTo for GLWESwitchingKey<D> {
116    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
117        writer.write_u64::<LittleEndian>(self.sk_in_n as u64)?;
118        writer.write_u64::<LittleEndian>(self.sk_out_n as u64)?;
119        self.key.write_to(writer)
120    }
121}
122
123#[derive(PartialEq, Eq)]
124pub struct GLWESwitchingKeyExec<D: Data, B: Backend> {
125    pub(crate) key: GGLWECiphertextExec<D, B>,
126    pub(crate) sk_in_n: usize,  // Degree of sk_in
127    pub(crate) sk_out_n: usize, // Degree of sk_out
128}
129
130impl<B: Backend> GLWESwitchingKeyExec<Vec<u8>, B> {
131    pub fn alloc(module: &Module<B>, basek: usize, k: usize, rows: usize, digits: usize, rank_in: usize, rank_out: usize) -> Self
132    where
133        Module<B>: GGLWEExecLayoutFamily<B>,
134    {
135        GLWESwitchingKeyExec::<Vec<u8>, B> {
136            key: GGLWECiphertextExec::alloc(module, basek, k, rows, digits, rank_in, rank_out),
137            sk_in_n: 0,
138            sk_out_n: 0,
139        }
140    }
141
142    pub fn bytes_of(
143        module: &Module<B>,
144        basek: usize,
145        k: usize,
146        rows: usize,
147        digits: usize,
148        rank_in: usize,
149        rank_out: usize,
150    ) -> usize
151    where
152        Module<B>: GGLWEExecLayoutFamily<B>,
153    {
154        GGLWECiphertextExec::bytes_of(module, basek, k, rows, digits, rank_in, rank_out)
155    }
156
157    pub fn from<DataOther: DataRef>(module: &Module<B>, other: &GLWESwitchingKey<DataOther>, scratch: &mut Scratch<B>) -> Self
158    where
159        Module<B>: GGLWEExecLayoutFamily<B>,
160    {
161        let mut ksk_exec: GLWESwitchingKeyExec<Vec<u8>, B> = Self::alloc(
162            module,
163            other.basek(),
164            other.k(),
165            other.rows(),
166            other.digits(),
167            other.rank_in(),
168            other.rank_out(),
169        );
170        ksk_exec.prepare(module, other, scratch);
171        ksk_exec
172    }
173}
174
175impl<D: Data, B: Backend> Infos for GLWESwitchingKeyExec<D, B> {
176    type Inner = VmpPMat<D, B>;
177
178    fn inner(&self) -> &Self::Inner {
179        self.key.inner()
180    }
181
182    fn basek(&self) -> usize {
183        self.key.basek()
184    }
185
186    fn k(&self) -> usize {
187        self.key.k()
188    }
189}
190
191impl<D: Data, B: Backend> GLWESwitchingKeyExec<D, B> {
192    pub fn rank(&self) -> usize {
193        self.key.data.cols_out() - 1
194    }
195
196    pub fn rank_in(&self) -> usize {
197        self.key.data.cols_in()
198    }
199
200    pub fn rank_out(&self) -> usize {
201        self.key.data.cols_out() - 1
202    }
203
204    pub fn digits(&self) -> usize {
205        self.key.digits()
206    }
207
208    pub fn sk_degree_in(&self) -> usize {
209        self.sk_in_n
210    }
211
212    pub fn sk_degree_out(&self) -> usize {
213        self.sk_out_n
214    }
215}
216
217impl<D: DataMut, B: Backend> GLWESwitchingKeyExec<D, B> {
218    pub fn prepare<DataOther>(&mut self, module: &Module<B>, other: &GLWESwitchingKey<DataOther>, scratch: &mut Scratch<B>)
219    where
220        DataOther: DataRef,
221        Module<B>: GGLWEExecLayoutFamily<B>,
222    {
223        self.key.prepare(module, &other.key, scratch);
224        self.sk_in_n = other.sk_in_n;
225        self.sk_out_n = other.sk_out_n;
226    }
227}