1use poulpy_hal::{
  2    api::{VmpPMatAlloc, VmpPMatAllocBytes},
  3    layouts::{Backend, Data, DataMut, DataRef, Module, Scratch},
  4};
  5
  6use crate::layouts::{
  7    Base2K, Degree, Dnum, Dsize, GGLWEInfos, GLWEInfos, GLWESwitchingKey, GLWESwitchingKeyToRef, LWEInfos, Rank, TorusPrecision,
  8    prepared::{GGLWEPrepare, GGLWEPrepareTmpBytes, GGLWEPrepared, GGLWEPreparedToMut, GGLWEPreparedToRef},
  9};
 10
 11#[derive(PartialEq, Eq)]
 12pub struct GLWESwitchingKeyPrepared<D: Data, B: Backend> {
 13    pub(crate) key: GGLWEPrepared<D, B>,
 14    pub(crate) sk_in_n: usize,  // Degree of sk_in
 15    pub(crate) sk_out_n: usize, // Degree of sk_out
 16}
 17
 18pub(crate) trait GLWESwitchingKeyPreparedSetMetaData {
 19    fn set_sk_in_n(&mut self, sk_in_n: usize);
 20    fn set_sk_out_n(&mut self, sk_out_n: usize);
 21}
 22
 23impl<D: DataMut, B: Backend> GLWESwitchingKeyPreparedSetMetaData for GLWESwitchingKeyPrepared<D, B> {
 24    fn set_sk_in_n(&mut self, sk_in_n: usize) {
 25        self.sk_in_n = sk_in_n
 26    }
 27
 28    fn set_sk_out_n(&mut self, sk_out_n: usize) {
 29        self.sk_out_n = self.sk_out_n
 30    }
 31}
 32
 33impl<D: Data, B: Backend> LWEInfos for GLWESwitchingKeyPrepared<D, B> {
 34    fn n(&self) -> Degree {
 35        self.key.n()
 36    }
 37
 38    fn base2k(&self) -> Base2K {
 39        self.key.base2k()
 40    }
 41
 42    fn k(&self) -> TorusPrecision {
 43        self.key.k()
 44    }
 45
 46    fn size(&self) -> usize {
 47        self.key.size()
 48    }
 49}
 50
 51impl<D: Data, B: Backend> GLWEInfos for GLWESwitchingKeyPrepared<D, B> {
 52    fn rank(&self) -> Rank {
 53        self.rank_out()
 54    }
 55}
 56
 57impl<D: Data, B: Backend> GGLWEInfos for GLWESwitchingKeyPrepared<D, B> {
 58    fn rank_in(&self) -> Rank {
 59        self.key.rank_in()
 60    }
 61
 62    fn rank_out(&self) -> Rank {
 63        self.key.rank_out()
 64    }
 65
 66    fn dsize(&self) -> Dsize {
 67        self.key.dsize()
 68    }
 69
 70    fn dnum(&self) -> Dnum {
 71        self.key.dnum()
 72    }
 73}
 74
 75impl<B: Backend> GLWESwitchingKeyPrepared<Vec<u8>, B> {
 76    pub fn alloc<A>(module: &Module<B>, infos: &A) -> Self
 77    where
 78        A: GGLWEInfos,
 79        Module<B>: VmpPMatAlloc<B>,
 80    {
 81        debug_assert_eq!(module.n() as u32, infos.n(), "module.n() != infos.n()");
 82        GLWESwitchingKeyPrepared::<Vec<u8>, B> {
 83            key: GGLWEPrepared::alloc(module, infos),
 84            sk_in_n: 0,
 85            sk_out_n: 0,
 86        }
 87    }
 88
 89    pub fn alloc_with(
 90        module: &Module<B>,
 91        base2k: Base2K,
 92        k: TorusPrecision,
 93        rank_in: Rank,
 94        rank_out: Rank,
 95        dnum: Dnum,
 96        dsize: Dsize,
 97    ) -> Self
 98    where
 99        Module<B>: VmpPMatAlloc<B>,
100    {
101        GLWESwitchingKeyPrepared::<Vec<u8>, B> {
102            key: GGLWEPrepared::alloc_with(module, base2k, k, rank_in, rank_out, dnum, dsize),
103            sk_in_n: 0,
104            sk_out_n: 0,
105        }
106    }
107
108    pub fn alloc_bytes<A>(module: &Module<B>, infos: &A) -> usize
109    where
110        A: GGLWEInfos,
111        Module<B>: VmpPMatAllocBytes,
112    {
113        debug_assert_eq!(module.n() as u32, infos.n(), "module.n() != infos.n()");
114        GGLWEPrepared::alloc_bytes(module, infos)
115    }
116
117    pub fn alloc_bytes_with(
118        module: &Module<B>,
119        base2k: Base2K,
120        k: TorusPrecision,
121        rank_in: Rank,
122        rank_out: Rank,
123        dnum: Dnum,
124        dsize: Dsize,
125    ) -> usize
126    where
127        Module<B>: VmpPMatAllocBytes,
128    {
129        GGLWEPrepared::alloc_bytes_with(module, base2k, k, rank_in, rank_out, dnum, dsize)
130    }
131}
132
133pub trait GLWESwitchingKeyPrepareTmpBytes {
134    fn glwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
135    where
136        A: GGLWEInfos;
137}
138
139impl<B: Backend> GLWESwitchingKeyPrepareTmpBytes for Module<B>
140where
141    Module<B>: GGLWEPrepareTmpBytes,
142{
143    fn glwe_switching_key_prepare_tmp_bytes<A>(&self, infos: &A) -> usize
144    where
145        A: GGLWEInfos,
146    {
147        self.gglwe_prepare_tmp_bytes(infos)
148    }
149}
150
151impl<B: Backend> GLWESwitchingKeyPrepared<Vec<u8>, B>
152where
153    Module<B>: GGLWEPrepareTmpBytes,
154{
155    pub fn prepare_tmp_bytes(&self, module: &Module<B>) -> usize {
156        module.gglwe_prepare_tmp_bytes(self)
157    }
158}
159
160pub trait GLWESwitchingKeyPrepare<B: Backend> {
161    fn glwe_switching_prepare<R, O>(&self, res: &R, other: &O, scratch: &mut Scratch<B>)
162    where
163        R: GLWESwitchingKeyPreparedToMut<B> + GLWESwitchingKeyPreparedSetMetaData,
164        O: GLWESwitchingKeyToRef;
165}
166
167impl<B: Backend> GLWESwitchingKeyPrepare<B> for Module<B>
168where
169    Module<B>: GGLWEPrepare<B>,
170{
171    fn glwe_switching_prepare<R, O>(&self, res: &R, other: &O, scratch: &mut Scratch<B>)
172    where
173        R: GLWESwitchingKeyPreparedToMut<B> + GLWESwitchingKeyPreparedSetMetaData,
174        O: GLWESwitchingKeyToRef,
175    {
176        self.gglwe_prepare(&res.to_mut(), other, scratch);
177        res.set_sk_in_n(other.sk_in_n);
178        res.set_sk_out_n(other.sk_out_n);
179    }
180}
181
182impl<D: DataMut, B: Backend> GLWESwitchingKeyPrepared<D, B> {
183    pub fn prepare<O>(&mut self, module: &Module<B>, other: &O, scratch: &mut Scratch<B>)
184    where
185        O: GLWESwitchingKeyToRef,
186        Module<B>: GLWESwitchingKeyPrepare<B>,
187    {
188        module.glwe_switching_prepare(self, other, scratch);
189    }
190}
191
192pub trait GLWESwitchingKeyPrepareAlloc<B: Backend> {
193    fn glwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
194    where
195        O: GLWESwitchingKeyToRef;
196}
197
198impl<B: Backend> GLWESwitchingKeyPrepareAlloc<B> for Module<B>
199where
200    Module<B>: GLWESwitchingKeyPrepare<B>,
201{
202    fn glwe_switching_key_prepare_alloc<O>(&self, other: &O, scratch: &mut Scratch<B>)
203    where
204        O: GLWESwitchingKeyToRef,
205    {
206        let mut ct_prepared: GLWESwitchingKeyPrepared<Vec<u8>, B> = GLWESwitchingKeyPrepared::alloc(self, self);
207        self.glwe_switching_prepare(&mut ct_prepared, other, scratch);
208        ct_prepared
209    }
210}
211
212impl<D: DataRef> GLWESwitchingKey<D> {
213    pub fn prepare_alloc<B: Backend>(&self, module: &Module<B>, scratch: &mut Scratch<B>) -> GLWESwitchingKeyPrepared<Vec<u8>, B>
214    where
215        Module<B>: GLWESwitchingKeyPrepareAlloc<B>,
216    {
217        module.glwe_switching_key_prepare_alloc(self, scratch);
218    }
219}
220
221pub trait GLWESwitchingKeyPreparedToMut<B: Backend> {
222    fn to_mut(&mut self) -> GLWESwitchingKeyPrepared<&mut [u8], B>;
223}
224
225impl<D: DataMut, B: Backend> GLWESwitchingKeyPreparedToMut<B> for GLWESwitchingKeyPrepared<D, B> {
226    fn to_mut(&mut self) -> GLWESwitchingKeyPrepared<&mut [u8], B> {
227        GLWESwitchingKeyPrepared {
228            sk_in_n: self.sk_in_n,
229            sk_out_n: self.sk_out_n,
230            key: self.key.to_mut(),
231        }
232    }
233}
234
235pub trait GLWESwitchingKeyPreparedToRef<B: Backend> {
236    fn to_ref(&self) -> GLWESwitchingKeyPrepared<&[u8], B>;
237}
238
239impl<D: DataRef, B: Backend> GLWESwitchingKeyPreparedToRef<B> for GLWESwitchingKeyPrepared<D, B> {
240    fn to_ref(&self) -> GLWESwitchingKeyPrepared<&[u8], B> {
241        GLWESwitchingKeyPrepared {
242            sk_in_n: self.sk_in_n,
243            sk_out_n: self.sk_out_n,
244            key: self.key.to_ref(),
245        }
246    }
247}