1use poulpy_hal::layouts::{Backend, Data, DataMut, DataRef, Module, Scratch};
  2
  3use crate::layouts::{
  4    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, GLWESwitchingKeySetMetaData, GLWESwitchingKeyToRef,
  5    GLWESwtichingKeyGetMetaData, LWEInfos, Rank, TorusPrecision,
  6    prepared::{
  7        GGLWEPrepare, GGLWEPrepareTmpBytes, GGLWEPrepared, GGLWEPreparedAlloc, GGLWEPreparedAllocBytes,
  8        GGLWEPreparedAllocBytesFromInfos, GGLWEPreparedAllocFromInfos, GGLWEPreparedToMut, GGLWEPreparedToRef,
  9    },
 10};
 11
 12#[derive(PartialEq, Eq)]
 13pub struct GLWESwitchingKeyPrepared<D: Data, B: Backend> {
 14    pub(crate) key: GGLWEPrepared<D, B>,
 15    pub(crate) sk_in_n: usize,  // Degree of sk_in
 16    pub(crate) sk_out_n: usize, // Degree of sk_out
 17}
 18
 19impl<D: DataMut, B: Backend> GLWESwitchingKeySetMetaData for GLWESwitchingKeyPrepared<D, B> {
 20    fn set_sk_in_n(&mut self, sk_in_n: usize) {
 21        self.sk_in_n = sk_in_n
 22    }
 23
 24    fn set_sk_out_n(&mut self, sk_out_n: usize) {
 25        self.sk_out_n = sk_out_n
 26    }
 27}
 28
 29impl<D: DataRef, B: Backend> GLWESwtichingKeyGetMetaData for GLWESwitchingKeyPrepared<D, B> {
 30    fn sk_in_n(&self) -> usize {
 31        self.sk_in_n
 32    }
 33
 34    fn sk_out_n(&self) -> usize {
 35        self.sk_out_n
 36    }
 37}
 38
 39impl<D: Data, B: Backend> LWEInfos for GLWESwitchingKeyPrepared<D, B> {
 40    fn n(&self) -> Degree {
 41        self.key.n()
 42    }
 43
 44    fn base2k(&self) -> Base2K {
 45        self.key.base2k()
 46    }
 47
 48    fn k(&self) -> TorusPrecision {
 49        self.key.k()
 50    }
 51
 52    fn size(&self) -> usize {
 53        self.key.size()
 54    }
 55}
 56
 57impl<D: Data, B: Backend> GLWEInfos for GLWESwitchingKeyPrepared<D, B> {
 58    fn rank(&self) -> Rank {
 59        self.rank_out()
 60    }
 61}
 62
 63impl<D: Data, B: Backend> GGLWEInfos for GLWESwitchingKeyPrepared<D, B> {
 64    fn rank_in(&self) -> Rank {
 65        self.key.rank_in()
 66    }
 67
 68    fn rank_out(&self) -> Rank {
 69        self.key.rank_out()
 70    }
 71
 72    fn dsize(&self) -> Dsize {
 73        self.key.dsize()
 74    }
 75
 76    fn dnum(&self) -> Dnum {
 77        self.key.dnum()
 78    }
 79}
 80
 81pub trait GLWESwitchingKeyPreparedAlloc<B: Backend> {
 82    fn glwe_switching_key_prepared_alloc(
 83        &self,
 84        base2k: Base2K,
 85        k: TorusPrecision,
 86        rank_in: Rank,
 87        rank_out: Rank,
 88        dnum: Dnum,
 89        dsize: Dsize,
 90    ) -> GLWESwitchingKeyPrepared<Vec<u8>, B>;
 91}
 92
 93impl<B: Backend> GLWESwitchingKeyPreparedAlloc<B> for Module<B>
 94where
 95    Module<B>: GGLWEPreparedAlloc<B>,
 96{
 97    fn glwe_switching_key_prepared_alloc(
 98        &self,
 99        base2k: Base2K,
100        k: TorusPrecision,
101        rank_in: Rank,
102        rank_out: Rank,
103        dnum: Dnum,
104        dsize: Dsize,
105    ) -> GLWESwitchingKeyPrepared<Vec<u8>, B> {
106        GLWESwitchingKeyPrepared::<Vec<u8>, B> {
107            key: self.gglwe_prepared_alloc(base2k, k, rank_in, rank_out, dnum, dsize),
108            sk_in_n: 0,
109            sk_out_n: 0,
110        }
111    }
112}
113
114pub trait GLWESwitchingKeyPreparedAllocFromInfos<B: Backend> {
115    fn glwe_switching_key_prepared_alloc_from_infos<A>(&self, infos: &A) -> GLWESwitchingKeyPrepared<Vec<u8>, B>
116    where
117        A: GGLWEInfos;
118}
119
120impl<B: Backend> GLWESwitchingKeyPreparedAllocFromInfos<B> for Module<B>
121where
122    Module<B>: GGLWEPreparedAllocFromInfos<B>,
123{
124    fn glwe_switching_key_prepared_alloc_from_infos<A>(&self, infos: &A) -> GLWESwitchingKeyPrepared<Vec<u8>, B>
125    where
126        A: GGLWEInfos,
127    {
128        GLWESwitchingKeyPrepared::<Vec<u8>, B> {
129            key: self.gglwe_prepared_alloc_from_infos(infos),
130            sk_in_n: 0,
131            sk_out_n: 0,
132        }
133    }
134}
135
136pub trait GLWESwitchingKeyPreparedAllocBytes<B: Backend> {
137    fn glwe_switching_key_prepared_alloc_bytes(
138        &self,
139        base2k: Base2K,
140        k: TorusPrecision,
141        rank_in: Rank,
142        rank_out: Rank,
143        dnum: Dnum,
144        dsize: Dsize,
145    ) -> usize;
146}
147
148impl<B: Backend> GLWESwitchingKeyPreparedAllocBytes<B> for Module<B>
149where
150    Module<B>: GGLWEPreparedAllocBytes<B>,
151{
152    fn glwe_switching_key_prepared_alloc_bytes(
153        &self,
154        base2k: Base2K,
155        k: TorusPrecision,
156        rank_in: Rank,
157        rank_out: Rank,
158        dnum: Dnum,
159        dsize: Dsize,
160    ) -> usize {
161        self.gglwe_prepared_alloc_bytes(base2k, k, rank_in, rank_out, dnum, dsize)
162    }
163}
164
165pub trait GLWESwitchingKeyPreparedAllocBytesFromInfos<B: Backend> {
166    fn glwe_switching_key_prepared_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
167    where
168        A: GGLWEInfos;
169}
170
171impl<B: Backend> GLWESwitchingKeyPreparedAllocBytesFromInfos<B> for Module<B>
172where
173    Module<B>: GGLWEPreparedAllocBytesFromInfos<B>,
174{
175    fn glwe_switching_key_prepared_alloc_bytes_from_infos<A>(&self, infos: &A) -> usize
176    where
177        A: GGLWEInfos,
178    {
179        self.gglwe_prepared_alloc_bytes_from_infos(infos)
180    }
181}
182
183impl<B: Backend> GLWESwitchingKeyPrepared<Vec<u8>, B> {
184    pub fn alloc_from_infos<A>(module: &Module<B>, infos: &A) -> Self
185    where
186        A: GGLWEInfos,
187        Module<B>: GLWESwitchingKeyPreparedAllocFromInfos<B>,
188    {
189        module.glwe_switching_key_prepared_alloc_from_infos(infos)
190    }
191
192    pub fn alloc(
193        module: &Module<B>,
194        base2k: Base2K,
195        k: TorusPrecision,
196        rank_in: Rank,
197        rank_out: Rank,
198        dnum: Dnum,
199        dsize: Dsize,
200    ) -> Self
201    where
202        Module<B>: GLWESwitchingKeyPreparedAlloc<B>,
203    {
204        module.glwe_switching_key_prepared_alloc(base2k, k, rank_in, rank_out, dnum, dsize)
205    }
206
207    pub fn alloc_bytes_from_infos<A>(module: &Module<B>, infos: &A) -> usize
208    where
209        A: GGLWEInfos,
210        Module<B>: GLWESwitchingKeyPreparedAllocBytesFromInfos<B>,
211    {
212        module.glwe_switching_key_prepared_alloc_bytes_from_infos(infos)
213    }
214
215    pub fn alloc_bytes(
216        module: &Module<B>,
217        base2k: Base2K,
218        k: TorusPrecision,
219        rank_in: Rank,
220        rank_out: Rank,
221        dnum: Dnum,
222        dsize: Dsize,
223    ) -> usize
224    where
225        Module<B>: GLWESwitchingKeyPreparedAllocBytes<B>,
226    {
227        module.glwe_switching_key_prepared_alloc_bytes(base2k, k, rank_in, rank_out, dnum, dsize)
228    }
229}
230
231pub trait GLWESwitchingKeyPrepareTmpBytes {
232    fn glwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
233    where
234        A: GGLWEInfos;
235}
236
237impl<B: Backend> GLWESwitchingKeyPrepareTmpBytes for Module<B>
238where
239    Module<B>: GGLWEPrepareTmpBytes,
240{
241    fn glwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
242    where
243        A: GGLWEInfos,
244    {
245        self.gglwe_prepare_tmp_bytes(infos)
246    }
247}
248
249impl<B: Backend> GLWESwitchingKeyPrepared<Vec<u8>, B>
250where
251    Module<B>: GGLWEPrepareTmpBytes,
252{
253    pub fn prepare_tmp_bytes(&self, module: &Module<B>) -> usize {
254        module.gglwe_prepare_tmp_bytes(self)
255    }
256}
257
258pub trait GLWESwitchingKeyPrepare<B: Backend> {
259    fn glwe_switching_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
260    where
261        R: GLWESwitchingKeyPreparedToMut<B> + GLWESwitchingKeySetMetaData,
262        O: GLWESwitchingKeyToRef + GLWESwtichingKeyGetMetaData;
263}
264
265impl<B: Backend> GLWESwitchingKeyPrepare<B> for Module<B>
266where
267    Module<B>: GGLWEPrepare<B>,
268{
269    fn glwe_switching_prepare<R, O>(&self, res: &mut R, other: &O, scratch: &mut Scratch<B>)
270    where
271        R: GLWESwitchingKeyPreparedToMut<B> + GLWESwitchingKeySetMetaData,
272        O: GLWESwitchingKeyToRef + GLWESwtichingKeyGetMetaData,
273    {
274        self.gglwe_prepare(&mut res.to_mut().key, &other.to_ref().key, scratch);
275        res.set_sk_in_n(other.sk_in_n());
276        res.set_sk_out_n(other.sk_out_n());
277    }
278}
279
280impl<D: DataMut, B: Backend> GLWESwitchingKeyPrepared<D, B> {
281    pub fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &mut Scratch<B>)
282    where
283        O: GLWESwitchingKeyToRef + GLWESwtichingKeyGetMetaData,
284        Module<B>: GLWESwitchingKeyPrepare<B>,
285    {
286        module.glwe_switching_prepare(self, other, scratch);
287    }
288}
289
290pub trait GLWESwitchingKeyPreparedToMut<B: Backend> {
291    fn to_mut(&mut self) -> GLWESwitchingKeyPrepared<&mut [u8], B>;
292}
293
294impl<D: DataMut, B: Backend> GLWESwitchingKeyPreparedToMut<B> for GLWESwitchingKeyPrepared<D, B> {
295    fn to_mut(&mut self) -> GLWESwitchingKeyPrepared<&mut [u8], B> {
296        GLWESwitchingKeyPrepared {
297            sk_in_n: self.sk_in_n,
298            sk_out_n: self.sk_out_n,
299            key: self.key.to_mut(),
300        }
301    }
302}
303
304pub trait GLWESwitchingKeyPreparedToRef<B: Backend> {
305    fn to_ref(&self) -> GLWESwitchingKeyPrepared<&[u8], B>;
306}
307
308impl<D: DataRef, B: Backend> GLWESwitchingKeyPreparedToRef<B> for GLWESwitchingKeyPrepared<D, B> {
309    fn to_ref(&self) -> GLWESwitchingKeyPrepared<&[u8], B> {
310        GLWESwitchingKeyPrepared {
311            sk_in_n: self.sk_in_n,
312            sk_out_n: self.sk_out_n,
313            key: self.key.to_ref(),
314        }
315    }
316}