1use backend::hal::{
  2    api::{
  3        ScratchOwnedAlloc, ScratchOwnedBorrow, VecZnxAddScalarInplace, VecZnxAutomorphismInplace, VecZnxEncodeCoeffsi64,
  4        VecZnxSwithcDegree, ZnxView,
  5    },
  6    layouts::{Backend, Module, ScratchOwned},
  7    oep::{
  8        ScratchAvailableImpl, ScratchOwnedAllocImpl, ScratchOwnedBorrowImpl, TakeScalarZnxImpl, TakeSvpPPolImpl,
  9        TakeVecZnxBigImpl, TakeVecZnxDftImpl, TakeVecZnxImpl,
 10    },
 11};
 12use sampling::source::Source;
 13
 14use crate::{
 15    GGLWEEncryptSkFamily, GGLWEExecLayoutFamily, GLWECiphertext, GLWEDecryptFamily, GLWEKeyswitchFamily, GLWEPlaintext,
 16    GLWESecret, GLWESecretExec, GLWEToLWESwitchingKey, GLWEToLWESwitchingKeyExec, Infos, LWECiphertext, LWEPlaintext, LWESecret,
 17    LWESwitchingKey, LWESwitchingKeyExec, LWEToGLWESwitchingKey, LWEToGLWESwitchingKeyExec,
 18};
 19
 20pub(crate) trait LWETestModuleFamily<B: Backend> = GGLWEEncryptSkFamily<B>
 21    + GLWEDecryptFamily<B>
 22    + VecZnxSwithcDegree
 23    + VecZnxAddScalarInplace
 24    + GGLWEExecLayoutFamily<B>
 25    + GLWEKeyswitchFamily<B>
 26    + VecZnxEncodeCoeffsi64
 27    + VecZnxAutomorphismInplace;
 28
 29pub(crate) trait LWETestScratchFamily<B: Backend> = TakeScalarZnxImpl<B>
 30    + TakeVecZnxDftImpl<B>
 31    + ScratchAvailableImpl<B>
 32    + TakeVecZnxImpl<B>
 33    + TakeVecZnxBigImpl<B>
 34    + TakeSvpPPolImpl<B>
 35    + ScratchOwnedAllocImpl<B>
 36    + ScratchOwnedBorrowImpl<B>;
 37
 38pub(crate) fn test_lwe_to_glwe<B: Backend>(module: &Module<B>)
 39where
 40    Module<B>: LWETestModuleFamily<B>,
 41    B: LWETestScratchFamily<B>,
 42{
 43    let n: usize = module.n();
 44    let basek: usize = 17;
 45    let sigma: f64 = 3.2;
 46
 47    let rank: usize = 2;
 48
 49    let n_lwe: usize = 22;
 50    let k_lwe_ct: usize = 2 * basek;
 51    let k_lwe_pt: usize = 8;
 52
 53    let k_glwe_ct: usize = 3 * basek;
 54
 55    let k_ksk: usize = k_lwe_ct + basek;
 56
 57    let mut source_xs: Source = Source::new([0u8; 32]);
 58    let mut source_xa: Source = Source::new([0u8; 32]);
 59    let mut source_xe: Source = Source::new([0u8; 32]);
 60
 61    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
 62        LWEToGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank)
 63            | GLWECiphertext::from_lwe_scratch_space(module, n, basek, k_lwe_ct, k_glwe_ct, k_ksk, rank)
 64            | GLWECiphertext::decrypt_scratch_space(module, n, basek, k_glwe_ct),
 65    );
 66
 67    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
 68    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
 69
 70    let sk_glwe_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_glwe);
 71
 72    let mut sk_lwe: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe);
 73    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
 74
 75    let data: i64 = 17;
 76
 77    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
 78    module.encode_coeff_i64(basek, &mut lwe_pt.data, 0, k_lwe_pt, 0, data, k_lwe_pt);
 79
 80    let mut lwe_ct: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
 81    lwe_ct.encrypt_sk(
 82        module,
 83        &lwe_pt,
 84        &sk_lwe,
 85        &mut source_xa,
 86        &mut source_xe,
 87        sigma,
 88    );
 89
 90    let mut ksk: LWEToGLWESwitchingKey<Vec<u8>> = LWEToGLWESwitchingKey::alloc(n, basek, k_ksk, lwe_ct.size(), rank);
 91
 92    ksk.encrypt_sk(
 93        module,
 94        &sk_lwe,
 95        &sk_glwe,
 96        &mut source_xa,
 97        &mut source_xe,
 98        sigma,
 99        scratch.borrow(),
100    );
101
102    let mut glwe_ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(n, basek, k_glwe_ct, rank);
103
104    let ksk_exec: LWEToGLWESwitchingKeyExec<Vec<u8>, B> = LWEToGLWESwitchingKeyExec::from(module, &ksk, scratch.borrow());
105
106    glwe_ct.from_lwe(module, &lwe_ct, &ksk_exec, scratch.borrow());
107
108    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_glwe_ct);
109    glwe_ct.decrypt(module, &mut glwe_pt, &sk_glwe_exec, scratch.borrow());
110
111    assert_eq!(glwe_pt.data.at(0, 0)[0], lwe_pt.data.at(0, 0)[0]);
112}
113
114pub(crate) fn test_glwe_to_lwe<B: Backend>(module: &Module<B>)
115where
116    Module<B>: LWETestModuleFamily<B>,
117    B: LWETestScratchFamily<B>,
118{
119    let n: usize = module.n();
120    let basek: usize = 17;
121    let sigma: f64 = 3.2;
122
123    let rank: usize = 2;
124
125    let n_lwe: usize = 22;
126    let k_lwe_ct: usize = 2 * basek;
127    let k_lwe_pt: usize = 8;
128
129    let k_glwe_ct: usize = 3 * basek;
130
131    let k_ksk: usize = k_lwe_ct + basek;
132
133    let mut source_xs: Source = Source::new([0u8; 32]);
134    let mut source_xa: Source = Source::new([0u8; 32]);
135    let mut source_xe: Source = Source::new([0u8; 32]);
136
137    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
138        LWEToGLWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk, rank)
139            | LWECiphertext::from_glwe_scratch_space(module, n, basek, k_lwe_ct, k_glwe_ct, k_ksk, rank)
140            | GLWECiphertext::decrypt_scratch_space(module, n, basek, k_glwe_ct),
141    );
142
143    let mut sk_glwe: GLWESecret<Vec<u8>> = GLWESecret::alloc(n, rank);
144    sk_glwe.fill_ternary_prob(0.5, &mut source_xs);
145
146    let sk_glwe_exec: GLWESecretExec<Vec<u8>, B> = GLWESecretExec::from(module, &sk_glwe);
147
148    let mut sk_lwe = LWESecret::alloc(n_lwe);
149    sk_lwe.fill_ternary_prob(0.5, &mut source_xs);
150
151    let data: i64 = 17;
152    let mut glwe_pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(n, basek, k_glwe_ct);
153    module.encode_coeff_i64(basek, &mut glwe_pt.data, 0, k_lwe_pt, 0, data, k_lwe_pt);
154
155    let mut glwe_ct = GLWECiphertext::alloc(n, basek, k_glwe_ct, rank);
156    glwe_ct.encrypt_sk(
157        module,
158        &glwe_pt,
159        &sk_glwe_exec,
160        &mut source_xa,
161        &mut source_xe,
162        sigma,
163        scratch.borrow(),
164    );
165
166    let mut ksk: GLWEToLWESwitchingKey<Vec<u8>> = GLWEToLWESwitchingKey::alloc(n, basek, k_ksk, glwe_ct.size(), rank);
167
168    ksk.encrypt_sk(
169        module,
170        &sk_lwe,
171        &sk_glwe,
172        &mut source_xa,
173        &mut source_xe,
174        sigma,
175        scratch.borrow(),
176    );
177
178    let mut lwe_ct: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe, basek, k_lwe_ct);
179
180    let ksk_exec: GLWEToLWESwitchingKeyExec<Vec<u8>, B> = GLWEToLWESwitchingKeyExec::from(module, &ksk, scratch.borrow());
181
182    lwe_ct.from_glwe(module, &glwe_ct, &ksk_exec, scratch.borrow());
183
184    let mut lwe_pt: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_ct);
185    lwe_ct.decrypt(module, &mut lwe_pt, &sk_lwe);
186
187    assert_eq!(glwe_pt.data.at(0, 0)[0], lwe_pt.data.at(0, 0)[0]);
188}
189
190pub(crate) fn test_keyswitch<B: Backend>(module: &Module<B>)
191where
192    Module<B>: LWETestModuleFamily<B>,
193    B: LWETestScratchFamily<B>,
194{
195    let n: usize = module.n();
196    let basek: usize = 17;
197    let sigma: f64 = 3.2;
198
199    let n_lwe_in: usize = 22;
200    let n_lwe_out: usize = 30;
201    let k_lwe_ct: usize = 2 * basek;
202    let k_lwe_pt: usize = 8;
203
204    let k_ksk: usize = k_lwe_ct + basek;
205
206    let mut source_xs: Source = Source::new([0u8; 32]);
207    let mut source_xa: Source = Source::new([0u8; 32]);
208    let mut source_xe: Source = Source::new([0u8; 32]);
209
210    let mut scratch: ScratchOwned<B> = ScratchOwned::alloc(
211        LWESwitchingKey::encrypt_sk_scratch_space(module, n, basek, k_ksk)
212            | LWECiphertext::keyswitch_scratch_space(module, n, basek, k_lwe_ct, k_lwe_ct, k_ksk),
213    );
214
215    let mut sk_lwe_in: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_in);
216    sk_lwe_in.fill_ternary_prob(0.5, &mut source_xs);
217
218    let mut sk_lwe_out: LWESecret<Vec<u8>> = LWESecret::alloc(n_lwe_out);
219    sk_lwe_out.fill_ternary_prob(0.5, &mut source_xs);
220
221    let data: i64 = 17;
222
223    let mut lwe_pt_in: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_pt);
224    module.encode_coeff_i64(basek, &mut lwe_pt_in.data, 0, k_lwe_pt, 0, data, k_lwe_pt);
225
226    let mut lwe_ct_in: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe_in, basek, k_lwe_ct);
227    lwe_ct_in.encrypt_sk(
228        module,
229        &lwe_pt_in,
230        &sk_lwe_in,
231        &mut source_xa,
232        &mut source_xe,
233        sigma,
234    );
235
236    let mut ksk: LWESwitchingKey<Vec<u8>> = LWESwitchingKey::alloc(n, basek, k_ksk, lwe_ct_in.size());
237
238    ksk.encrypt_sk(
239        module,
240        &sk_lwe_in,
241        &sk_lwe_out,
242        &mut source_xa,
243        &mut source_xe,
244        sigma,
245        scratch.borrow(),
246    );
247
248    let mut lwe_ct_out: LWECiphertext<Vec<u8>> = LWECiphertext::alloc(n_lwe_out, basek, k_lwe_ct);
249
250    let ksk_exec: LWESwitchingKeyExec<Vec<u8>, B> = LWESwitchingKeyExec::from(module, &ksk, scratch.borrow());
251
252    lwe_ct_out.keyswitch(module, &lwe_ct_in, &ksk_exec, scratch.borrow());
253
254    let mut lwe_pt_out: LWEPlaintext<Vec<u8>> = LWEPlaintext::alloc(basek, k_lwe_ct);
255    lwe_ct_out.decrypt(module, &mut lwe_pt_out, &sk_lwe_out);
256
257    assert_eq!(lwe_pt_in.data.at(0, 0)[0], lwe_pt_out.data.at(0, 0)[0]);
258}