1use crate::{
  2    elem::Infos,
  3    ggsw_ciphertext::GGSWCiphertext,
  4    glwe_ciphertext::GLWECiphertext,
  5    glwe_ciphertext_fourier::GLWECiphertextFourier,
  6    glwe_plaintext::GLWEPlaintext,
  7    keys::{SecretKey, SecretKeyFourier},
  8    keyswitch_key::GLWESwitchingKey,
  9    test_fft64::{gglwe::noise_grlwe_rlwe_product, ggsw::noise_rgsw_product},
 10};
 11use base2k::{FFT64, FillUniform, Module, ScalarZnx, ScalarZnxAlloc, ScratchOwned, Stats, VecZnxOps, VecZnxToMut, ZnxViewMut};
 12use sampling::source::Source;
 13
 14#[test]
 15fn keyswitch() {
 16    let module: Module<FFT64> = Module::<FFT64>::new(2048);
 17    let log_base2k: usize = 12;
 18    let log_k_grlwe: usize = 60;
 19    let log_k_rlwe_in: usize = 45;
 20    let log_k_rlwe_out: usize = 60;
 21    let rows: usize = (log_k_rlwe_in + log_base2k - 1) / log_base2k;
 22
 23    let rank: usize = 1;
 24
 25    let sigma: f64 = 3.2;
 26
 27    let mut ct_grlwe: GLWESwitchingKey<Vec<u8>, FFT64> =
 28        GLWESwitchingKey::new(&module, log_base2k, log_k_grlwe, rows, rank, rank);
 29    let mut ct_rlwe_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, log_base2k, log_k_rlwe_in, rank);
 30    let mut ct_rlwe_in_dft: GLWECiphertextFourier<Vec<u8>, FFT64> =
 31        GLWECiphertextFourier::new(&module, log_base2k, log_k_rlwe_in, rank);
 32    let mut ct_rlwe_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, log_base2k, log_k_rlwe_out, rank);
 33    let mut ct_rlwe_out_dft: GLWECiphertextFourier<Vec<u8>, FFT64> =
 34        GLWECiphertextFourier::new(&module, log_base2k, log_k_rlwe_out, rank);
 35    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_rlwe_in);
 36    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_rlwe_out);
 37
 38    let mut source_xs: Source = Source::new([0u8; 32]);
 39    let mut source_xe: Source = Source::new([0u8; 32]);
 40    let mut source_xa: Source = Source::new([0u8; 32]);
 41
 42    // Random input plaintext
 43    pt_want
 44        .data
 45        .fill_uniform(log_base2k, 0, pt_want.size(), &mut source_xa);
 46
 47    let mut scratch: ScratchOwned = ScratchOwned::new(
 48        GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ct_grlwe.size())
 49            | GLWECiphertext::decrypt_scratch_space(&module, ct_rlwe_out.size())
 50            | GLWECiphertext::encrypt_sk_scratch_space(&module, rank, ct_rlwe_in.size())
 51            | GLWECiphertextFourier::keyswitch_scratch_space(
 52                &module,
 53                ct_rlwe_out.size(),
 54                ct_rlwe_in.size(),
 55                ct_grlwe.size(),
 56            ),
 57    );
 58
 59    let mut sk0: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
 60    sk0.fill_ternary_prob(0.5, &mut source_xs);
 61
 62    let mut sk0_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
 63    sk0_dft.dft(&module, &sk0);
 64
 65    let mut sk1: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
 66    sk1.fill_ternary_prob(0.5, &mut source_xs);
 67
 68    let mut sk1_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
 69    sk1_dft.dft(&module, &sk1);
 70
 71    ct_grlwe.encrypt_sk(
 72        &module,
 73        &sk0.data,
 74        &sk1_dft,
 75        &mut source_xa,
 76        &mut source_xe,
 77        sigma,
 78        scratch.borrow(),
 79    );
 80
 81    ct_rlwe_in.encrypt_sk(
 82        &module,
 83        &pt_want,
 84        &sk0_dft,
 85        &mut source_xa,
 86        &mut source_xe,
 87        sigma,
 88        scratch.borrow(),
 89    );
 90
 91    ct_rlwe_in.dft(&module, &mut ct_rlwe_in_dft);
 92    ct_rlwe_out_dft.keyswitch(&module, &ct_rlwe_in_dft, &ct_grlwe, scratch.borrow());
 93    ct_rlwe_out_dft.idft(&module, &mut ct_rlwe_out, scratch.borrow());
 94
 95    ct_rlwe_out.decrypt(&module, &mut pt_have, &sk1_dft, scratch.borrow());
 96
 97    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
 98
 99    let noise_have: f64 = pt_have.data.std(0, log_base2k).log2();
100    let noise_want: f64 = noise_grlwe_rlwe_product(
101        module.n() as f64,
102        log_base2k,
103        0.5,
104        0.5,
105        0f64,
106        sigma * sigma,
107        0f64,
108        log_k_rlwe_in,
109        log_k_grlwe,
110    );
111
112    assert!(
113        (noise_have - noise_want).abs() <= 0.1,
114        "{} {}",
115        noise_have,
116        noise_want
117    );
118}
119
120#[test]
121fn keyswich_inplace() {
122    let module: Module<FFT64> = Module::<FFT64>::new(2048);
123    let log_base2k: usize = 12;
124    let log_k_grlwe: usize = 60;
125    let log_k_rlwe: usize = 45;
126    let rows: usize = (log_k_rlwe + log_base2k - 1) / log_base2k;
127    let rank: usize = 1;
128
129    let sigma: f64 = 3.2;
130
131    let mut ct_grlwe: GLWESwitchingKey<Vec<u8>, FFT64> =
132        GLWESwitchingKey::new(&module, log_base2k, log_k_grlwe, rows, rank, rank);
133    let mut ct_rlwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, log_base2k, log_k_rlwe, rank);
134    let mut ct_rlwe_dft: GLWECiphertextFourier<Vec<u8>, FFT64> =
135        GLWECiphertextFourier::new(&module, log_base2k, log_k_rlwe, rank);
136    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_rlwe);
137    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_rlwe);
138
139    let mut source_xs: Source = Source::new([0u8; 32]);
140    let mut source_xe: Source = Source::new([0u8; 32]);
141    let mut source_xa: Source = Source::new([0u8; 32]);
142
143    // Random input plaintext
144    pt_want
145        .data
146        .fill_uniform(log_base2k, 0, pt_want.size(), &mut source_xa);
147
148    let mut scratch: ScratchOwned = ScratchOwned::new(
149        GLWESwitchingKey::encrypt_sk_scratch_space(&module, rank, ct_grlwe.size())
150            | GLWECiphertext::decrypt_scratch_space(&module, ct_rlwe.size())
151            | GLWECiphertext::encrypt_sk_scratch_space(&module, rank, ct_rlwe.size())
152            | GLWECiphertextFourier::keyswitch_inplace_scratch_space(&module, ct_rlwe_dft.size(), ct_grlwe.size()),
153    );
154
155    let mut sk0: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
156    sk0.fill_ternary_prob(0.5, &mut source_xs);
157
158    let mut sk0_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
159    sk0_dft.dft(&module, &sk0);
160
161    let mut sk1: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
162    sk1.fill_ternary_prob(0.5, &mut source_xs);
163
164    let mut sk1_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
165    sk1_dft.dft(&module, &sk1);
166
167    ct_grlwe.encrypt_sk(
168        &module,
169        &sk0.data,
170        &sk1_dft,
171        &mut source_xa,
172        &mut source_xe,
173        sigma,
174        scratch.borrow(),
175    );
176
177    ct_rlwe.encrypt_sk(
178        &module,
179        &pt_want,
180        &sk0_dft,
181        &mut source_xa,
182        &mut source_xe,
183        sigma,
184        scratch.borrow(),
185    );
186
187    ct_rlwe.dft(&module, &mut ct_rlwe_dft);
188    ct_rlwe_dft.keyswitch_inplace(&module, &ct_grlwe, scratch.borrow());
189    ct_rlwe_dft.idft(&module, &mut ct_rlwe, scratch.borrow());
190
191    ct_rlwe.decrypt(&module, &mut pt_have, &sk1_dft, scratch.borrow());
192
193    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
194
195    let noise_have: f64 = pt_have.data.std(0, log_base2k).log2();
196    let noise_want: f64 = noise_grlwe_rlwe_product(
197        module.n() as f64,
198        log_base2k,
199        0.5,
200        0.5,
201        0f64,
202        sigma * sigma,
203        0f64,
204        log_k_rlwe,
205        log_k_grlwe,
206    );
207
208    assert!(
209        (noise_have - noise_want).abs() <= 0.1,
210        "{} {}",
211        noise_have,
212        noise_want
213    );
214}
215
216#[test]
217fn external_product() {
218    let module: Module<FFT64> = Module::<FFT64>::new(2048);
219    let log_base2k: usize = 12;
220    let log_k_grlwe: usize = 60;
221    let log_k_rlwe_in: usize = 45;
222    let log_k_rlwe_out: usize = 60;
223    let rows: usize = (log_k_rlwe_in + log_base2k - 1) / log_base2k;
224    let rank: usize = 1;
225
226    let sigma: f64 = 3.2;
227
228    let mut ct_rgsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, log_base2k, log_k_grlwe, rows, rank);
229    let mut ct_rlwe_in: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, log_base2k, log_k_rlwe_in, rank);
230    let mut ct_rlwe_out: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, log_base2k, log_k_rlwe_out, rank);
231    let mut ct_rlwe_dft_in: GLWECiphertextFourier<Vec<u8>, FFT64> =
232        GLWECiphertextFourier::new(&module, log_base2k, log_k_rlwe_in, rank);
233    let mut ct_rlwe_dft_out: GLWECiphertextFourier<Vec<u8>, FFT64> =
234        GLWECiphertextFourier::new(&module, log_base2k, log_k_rlwe_out, rank);
235    let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
236    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_rlwe_in);
237    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_rlwe_out);
238
239    let mut source_xs: Source = Source::new([0u8; 32]);
240    let mut source_xe: Source = Source::new([0u8; 32]);
241    let mut source_xa: Source = Source::new([0u8; 32]);
242
243    // Random input plaintext
244    pt_want
245        .data
246        .fill_uniform(log_base2k, 0, pt_want.size(), &mut source_xa);
247
248    pt_want.to_mut().at_mut(0, 0)[1] = 1;
249
250    let k: usize = 1;
251
252    pt_rgsw.raw_mut()[k] = 1; // X^{k}
253
254    let mut scratch: ScratchOwned = ScratchOwned::new(
255        GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct_rgsw.size())
256            | GLWECiphertext::decrypt_scratch_space(&module, ct_rlwe_out.size())
257            | GLWECiphertext::encrypt_sk_scratch_space(&module, rank, ct_rlwe_in.size())
258            | GLWECiphertext::external_product_scratch_space(
259                &module,
260                ct_rlwe_out.size(),
261                ct_rlwe_in.size(),
262                ct_rgsw.size(),
263            ),
264    );
265
266    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
267    sk.fill_ternary_prob(0.5, &mut source_xs);
268
269    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
270    sk_dft.dft(&module, &sk);
271
272    ct_rgsw.encrypt_sk(
273        &module,
274        &pt_rgsw,
275        &sk_dft,
276        &mut source_xa,
277        &mut source_xe,
278        sigma,
279        scratch.borrow(),
280    );
281
282    ct_rlwe_in.encrypt_sk(
283        &module,
284        &pt_want,
285        &sk_dft,
286        &mut source_xa,
287        &mut source_xe,
288        sigma,
289        scratch.borrow(),
290    );
291
292    ct_rlwe_in.dft(&module, &mut ct_rlwe_dft_in);
293    ct_rlwe_dft_out.external_product(&module, &ct_rlwe_dft_in, &ct_rgsw, scratch.borrow());
294    ct_rlwe_dft_out.idft(&module, &mut ct_rlwe_out, scratch.borrow());
295
296    ct_rlwe_out.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
297
298    module.vec_znx_rotate_inplace(k as i64, &mut pt_want, 0);
299
300    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
301
302    let noise_have: f64 = pt_have.data.std(0, log_base2k).log2();
303
304    let var_gct_err_lhs: f64 = sigma * sigma;
305    let var_gct_err_rhs: f64 = 0f64;
306
307    let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
308    let var_a0_err: f64 = sigma * sigma;
309    let var_a1_err: f64 = 1f64 / 12f64;
310
311    let noise_want: f64 = noise_rgsw_product(
312        module.n() as f64,
313        log_base2k,
314        0.5,
315        var_msg,
316        var_a0_err,
317        var_a1_err,
318        var_gct_err_lhs,
319        var_gct_err_rhs,
320        log_k_rlwe_in,
321        log_k_grlwe,
322    );
323
324    assert!(
325        (noise_have - noise_want).abs() <= 0.1,
326        "{} {}",
327        noise_have,
328        noise_want
329    );
330}
331
332#[test]
333fn external_product_inplace() {
334    let module: Module<FFT64> = Module::<FFT64>::new(2048);
335    let log_base2k: usize = 12;
336    let log_k_grlwe: usize = 60;
337    let log_k_rlwe_in: usize = 45;
338    let log_k_rlwe_out: usize = 60;
339    let rows: usize = (log_k_rlwe_in + log_base2k - 1) / log_base2k;
340    let rank: usize = 1;
341
342    let sigma: f64 = 3.2;
343
344    let mut ct_rgsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, log_base2k, log_k_grlwe, rows, rank);
345    let mut ct_rlwe: GLWECiphertext<Vec<u8>> = GLWECiphertext::new(&module, log_base2k, log_k_rlwe_in, rank);
346    let mut ct_rlwe_dft: GLWECiphertextFourier<Vec<u8>, FFT64> =
347        GLWECiphertextFourier::new(&module, log_base2k, log_k_rlwe_in, rank);
348    let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
349    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_rlwe_in);
350    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_rlwe_out);
351
352    let mut source_xs: Source = Source::new([0u8; 32]);
353    let mut source_xe: Source = Source::new([0u8; 32]);
354    let mut source_xa: Source = Source::new([0u8; 32]);
355
356    // Random input plaintext
357    pt_want
358        .data
359        .fill_uniform(log_base2k, 0, pt_want.size(), &mut source_xa);
360
361    pt_want.to_mut().at_mut(0, 0)[1] = 1;
362
363    let k: usize = 1;
364
365    pt_rgsw.raw_mut()[k] = 1; // X^{k}
366
367    let mut scratch: ScratchOwned = ScratchOwned::new(
368        GGSWCiphertext::encrypt_sk_scratch_space(&module, rank, ct_rgsw.size())
369            | GLWECiphertext::decrypt_scratch_space(&module, ct_rlwe.size())
370            | GLWECiphertext::encrypt_sk_scratch_space(&module, rank, ct_rlwe.size())
371            | GLWECiphertext::external_product_inplace_scratch_space(&module, ct_rlwe.size(), ct_rgsw.size()),
372    );
373
374    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
375    sk.fill_ternary_prob(0.5, &mut source_xs);
376
377    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
378    sk_dft.dft(&module, &sk);
379
380    ct_rgsw.encrypt_sk(
381        &module,
382        &pt_rgsw,
383        &sk_dft,
384        &mut source_xa,
385        &mut source_xe,
386        sigma,
387        scratch.borrow(),
388    );
389
390    ct_rlwe.encrypt_sk(
391        &module,
392        &pt_want,
393        &sk_dft,
394        &mut source_xa,
395        &mut source_xe,
396        sigma,
397        scratch.borrow(),
398    );
399
400    ct_rlwe.dft(&module, &mut ct_rlwe_dft);
401    ct_rlwe_dft.external_product_inplace(&module, &ct_rgsw, scratch.borrow());
402    ct_rlwe_dft.idft(&module, &mut ct_rlwe, scratch.borrow());
403
404    ct_rlwe.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
405
406    module.vec_znx_rotate_inplace(k as i64, &mut pt_want, 0);
407
408    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
409
410    let noise_have: f64 = pt_have.data.std(0, log_base2k).log2();
411
412    let var_gct_err_lhs: f64 = sigma * sigma;
413    let var_gct_err_rhs: f64 = 0f64;
414
415    let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
416    let var_a0_err: f64 = sigma * sigma;
417    let var_a1_err: f64 = 1f64 / 12f64;
418
419    let noise_want: f64 = noise_rgsw_product(
420        module.n() as f64,
421        log_base2k,
422        0.5,
423        var_msg,
424        var_a0_err,
425        var_a1_err,
426        var_gct_err_lhs,
427        var_gct_err_rhs,
428        log_k_rlwe_in,
429        log_k_grlwe,
430    );
431
432    assert!(
433        (noise_have - noise_want).abs() <= 0.1,
434        "{} {}",
435        noise_have,
436        noise_want
437    );
438}