1use base2k::{FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScratchOwned, Stats, VecZnxOps, ZnxViewMut};
  2use sampling::source::Source;
  3
  4use crate::{
  5    elem::{GetRow, Infos},
  6    external_product::{
  7        ExternalProduct, ExternalProductInplace, ExternalProductInplaceScratchSpace, ExternalProductScratchSpace,
  8    },
  9    ggsw::GGSWCiphertext,
 10    glwe::{GLWECiphertextFourier, GLWEPlaintext},
 11    keys::{SecretKey, SecretKeyFourier},
 12    keyswitch::{KeySwitch, KeySwitchInplace, KeySwitchInplaceScratchSpace, KeySwitchScratchSpace},
 13    keyswitch_key::GLWEKeySwitchKey,
 14    test_fft64::rgsw::noise_rgsw_product,
 15};
 16
 17#[test]
 18fn encrypt_sk() {
 19    let module: Module<FFT64> = Module::<FFT64>::new(2048);
 20    let log_base2k: usize = 8;
 21    let log_k_ct: usize = 54;
 22    let rows: usize = 4;
 23
 24    let sigma: f64 = 3.2;
 25    let bound: f64 = sigma * 6.0;
 26
 27    let mut ct: GLWEKeySwitchKey<Vec<u8>, FFT64> = GLWEKeySwitchKey::new(&module, log_base2k, log_k_ct, rows);
 28    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_ct);
 29    let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
 30
 31    let mut source_xs: Source = Source::new([0u8; 32]);
 32    let mut source_xe: Source = Source::new([0u8; 32]);
 33    let mut source_xa: Source = Source::new([0u8; 32]);
 34
 35    pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
 36
 37    let mut scratch: ScratchOwned = ScratchOwned::new(
 38        GLWEKeySwitchKey::encrypt_sk_scratch_space(&module, ct.size())
 39            | GLWECiphertextFourier::decrypt_scratch_space(&module, ct.size()),
 40    );
 41
 42    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
 43    sk.fill_ternary_prob(0.5, &mut source_xs);
 44
 45    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
 46    sk_dft.dft(&module, &sk);
 47
 48    ct.encrypt_sk(
 49        &module,
 50        &pt_scalar,
 51        &sk_dft,
 52        &mut source_xa,
 53        &mut source_xe,
 54        sigma,
 55        bound,
 56        scratch.borrow(),
 57    );
 58
 59    let mut ct_rlwe_dft: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, log_base2k, log_k_ct);
 60
 61    (0..ct.rows()).for_each(|row_i| {
 62        ct.get_row(&module, row_i, 0, &mut ct_rlwe_dft);
 63        ct_rlwe_dft.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
 64        module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &pt_scalar, 0);
 65        let std_pt: f64 = pt.data.std(0, log_base2k) * (log_k_ct as f64).exp2();
 66        assert!((sigma - std_pt).abs() <= 0.2, "{} {}", sigma, std_pt);
 67    });
 68}
 69
 70#[test]
 71fn keyswitch() {
 72    let module: Module<FFT64> = Module::<FFT64>::new(2048);
 73    let log_base2k: usize = 12;
 74    let log_k_grlwe: usize = 60;
 75    let rows: usize = (log_k_grlwe + log_base2k - 1) / log_base2k;
 76
 77    let sigma: f64 = 3.2;
 78    let bound: f64 = sigma * 6.0;
 79
 80    let mut ct_grlwe_s0s1: GLWEKeySwitchKey<Vec<u8>, FFT64> = GLWEKeySwitchKey::new(&module, log_base2k, log_k_grlwe, rows);
 81    let mut ct_grlwe_s1s2: GLWEKeySwitchKey<Vec<u8>, FFT64> = GLWEKeySwitchKey::new(&module, log_base2k, log_k_grlwe, rows);
 82    let mut ct_grlwe_s0s2: GLWEKeySwitchKey<Vec<u8>, FFT64> = GLWEKeySwitchKey::new(&module, log_base2k, log_k_grlwe, rows);
 83
 84    let mut source_xs: Source = Source::new([0u8; 32]);
 85    let mut source_xe: Source = Source::new([0u8; 32]);
 86    let mut source_xa: Source = Source::new([0u8; 32]);
 87
 88    let mut scratch: ScratchOwned = ScratchOwned::new(
 89        GLWEKeySwitchKey::encrypt_sk_scratch_space(&module, ct_grlwe_s0s1.size())
 90            | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_grlwe_s0s2.size())
 91            | GLWEKeySwitchKey::keyswitch_scratch_space(
 92                &module,
 93                ct_grlwe_s0s2.size(),
 94                ct_grlwe_s0s1.size(),
 95                ct_grlwe_s1s2.size(),
 96            ),
 97    );
 98
 99    let mut sk0: SecretKey<Vec<u8>> = SecretKey::new(&module);
100    sk0.fill_ternary_prob(0.5, &mut source_xs);
101
102    let mut sk0_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
103    sk0_dft.dft(&module, &sk0);
104
105    let mut sk1: SecretKey<Vec<u8>> = SecretKey::new(&module);
106    sk1.fill_ternary_prob(0.5, &mut source_xs);
107
108    let mut sk1_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
109    sk1_dft.dft(&module, &sk1);
110
111    let mut sk2: SecretKey<Vec<u8>> = SecretKey::new(&module);
112    sk2.fill_ternary_prob(0.5, &mut source_xs);
113
114    let mut sk2_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
115    sk2_dft.dft(&module, &sk2);
116
117    // GRLWE_{s1}(s0) = s0 -> s1
118    ct_grlwe_s0s1.encrypt_sk(
119        &module,
120        &sk0.data,
121        &sk1_dft,
122        &mut source_xa,
123        &mut source_xe,
124        sigma,
125        bound,
126        scratch.borrow(),
127    );
128
129    // GRLWE_{s2}(s1) -> s1 -> s2
130    ct_grlwe_s1s2.encrypt_sk(
131        &module,
132        &sk1.data,
133        &sk2_dft,
134        &mut source_xa,
135        &mut source_xe,
136        sigma,
137        bound,
138        scratch.borrow(),
139    );
140
141    // GRLWE_{s1}(s0) (x) GRLWE_{s2}(s1) = GRLWE_{s2}(s0)
142    ct_grlwe_s0s2.keyswitch(&module, &ct_grlwe_s0s1, &ct_grlwe_s1s2, scratch.borrow());
143
144    let mut ct_rlwe_dft_s0s2: GLWECiphertextFourier<Vec<u8>, FFT64> =
145        GLWECiphertextFourier::new(&module, log_base2k, log_k_grlwe);
146    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_grlwe);
147
148    (0..ct_grlwe_s0s2.rows()).for_each(|row_i| {
149        ct_grlwe_s0s2.get_row(&module, row_i, 0, &mut ct_rlwe_dft_s0s2);
150        ct_rlwe_dft_s0s2.decrypt(&module, &mut pt, &sk2_dft, scratch.borrow());
151        module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk0, 0);
152
153        let noise_have: f64 = pt.data.std(0, log_base2k).log2();
154        let noise_want: f64 = noise_grlwe_rlwe_product(
155            module.n() as f64,
156            log_base2k,
157            0.5,
158            0.5,
159            0f64,
160            sigma * sigma,
161            0f64,
162            log_k_grlwe,
163            log_k_grlwe,
164        );
165
166        assert!(
167            (noise_have - noise_want).abs() <= 0.1,
168            "{} {}",
169            noise_have,
170            noise_want
171        );
172    });
173}
174
175#[test]
176fn keyswitch_inplace() {
177    let module: Module<FFT64> = Module::<FFT64>::new(2048);
178    let log_base2k: usize = 12;
179    let log_k_grlwe: usize = 60;
180    let rows: usize = (log_k_grlwe + log_base2k - 1) / log_base2k;
181
182    let sigma: f64 = 3.2;
183    let bound: f64 = sigma * 6.0;
184
185    let mut ct_grlwe_s0s1: GLWEKeySwitchKey<Vec<u8>, FFT64> = GLWEKeySwitchKey::new(&module, log_base2k, log_k_grlwe, rows);
186    let mut ct_grlwe_s1s2: GLWEKeySwitchKey<Vec<u8>, FFT64> = GLWEKeySwitchKey::new(&module, log_base2k, log_k_grlwe, rows);
187
188    let mut source_xs: Source = Source::new([0u8; 32]);
189    let mut source_xe: Source = Source::new([0u8; 32]);
190    let mut source_xa: Source = Source::new([0u8; 32]);
191
192    let mut scratch: ScratchOwned = ScratchOwned::new(
193        GLWEKeySwitchKey::encrypt_sk_scratch_space(&module, ct_grlwe_s0s1.size())
194            | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_grlwe_s0s1.size())
195            | GLWEKeySwitchKey::keyswitch_inplace_scratch_space(&module, ct_grlwe_s0s1.size(), ct_grlwe_s1s2.size()),
196    );
197
198    let mut sk0: SecretKey<Vec<u8>> = SecretKey::new(&module);
199    sk0.fill_ternary_prob(0.5, &mut source_xs);
200
201    let mut sk0_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
202    sk0_dft.dft(&module, &sk0);
203
204    let mut sk1: SecretKey<Vec<u8>> = SecretKey::new(&module);
205    sk1.fill_ternary_prob(0.5, &mut source_xs);
206
207    let mut sk1_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
208    sk1_dft.dft(&module, &sk1);
209
210    let mut sk2: SecretKey<Vec<u8>> = SecretKey::new(&module);
211    sk2.fill_ternary_prob(0.5, &mut source_xs);
212
213    let mut sk2_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
214    sk2_dft.dft(&module, &sk2);
215
216    // GRLWE_{s1}(s0) = s0 -> s1
217    ct_grlwe_s0s1.encrypt_sk(
218        &module,
219        &sk0.data,
220        &sk1_dft,
221        &mut source_xa,
222        &mut source_xe,
223        sigma,
224        bound,
225        scratch.borrow(),
226    );
227
228    // GRLWE_{s2}(s1) -> s1 -> s2
229    ct_grlwe_s1s2.encrypt_sk(
230        &module,
231        &sk1.data,
232        &sk2_dft,
233        &mut source_xa,
234        &mut source_xe,
235        sigma,
236        bound,
237        scratch.borrow(),
238    );
239
240    // GRLWE_{s1}(s0) (x) GRLWE_{s2}(s1) = GRLWE_{s2}(s0)
241    ct_grlwe_s0s1.keyswitch_inplace(&module, &ct_grlwe_s1s2, scratch.borrow());
242
243    let ct_grlwe_s0s2: GLWEKeySwitchKey<Vec<u8>, FFT64> = ct_grlwe_s0s1;
244
245    let mut ct_rlwe_dft_s0s2: GLWECiphertextFourier<Vec<u8>, FFT64> =
246        GLWECiphertextFourier::new(&module, log_base2k, log_k_grlwe);
247    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_grlwe);
248
249    (0..ct_grlwe_s0s2.rows()).for_each(|row_i| {
250        ct_grlwe_s0s2.get_row(&module, row_i, 0, &mut ct_rlwe_dft_s0s2);
251        ct_rlwe_dft_s0s2.decrypt(&module, &mut pt, &sk2_dft, scratch.borrow());
252        module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk0, 0);
253
254        let noise_have: f64 = pt.data.std(0, log_base2k).log2();
255        let noise_want: f64 = noise_grlwe_rlwe_product(
256            module.n() as f64,
257            log_base2k,
258            0.5,
259            0.5,
260            0f64,
261            sigma * sigma,
262            0f64,
263            log_k_grlwe,
264            log_k_grlwe,
265        );
266
267        assert!(
268            (noise_have - noise_want).abs() <= 0.1,
269            "{} {}",
270            noise_have,
271            noise_want
272        );
273    });
274}
275
276#[test]
277fn external_product() {
278    let module: Module<FFT64> = Module::<FFT64>::new(2048);
279    let log_base2k: usize = 12;
280    let log_k_grlwe: usize = 60;
281    let rows: usize = (log_k_grlwe + log_base2k - 1) / log_base2k;
282
283    let sigma: f64 = 3.2;
284    let bound: f64 = sigma * 6.0;
285
286    let mut ct_grlwe_in: GLWEKeySwitchKey<Vec<u8>, FFT64> = GLWEKeySwitchKey::new(&module, log_base2k, log_k_grlwe, rows);
287    let mut ct_grlwe_out: GLWEKeySwitchKey<Vec<u8>, FFT64> = GLWEKeySwitchKey::new(&module, log_base2k, log_k_grlwe, rows);
288    let mut ct_rgsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, log_base2k, log_k_grlwe, rows);
289
290    let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
291    let mut pt_grlwe: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
292
293    let mut source_xs: Source = Source::new([0u8; 32]);
294    let mut source_xe: Source = Source::new([0u8; 32]);
295    let mut source_xa: Source = Source::new([0u8; 32]);
296
297    let mut scratch: ScratchOwned = ScratchOwned::new(
298        GLWEKeySwitchKey::encrypt_sk_scratch_space(&module, ct_grlwe_in.size())
299            | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_grlwe_out.size())
300            | GLWEKeySwitchKey::external_product_scratch_space(
301                &module,
302                ct_grlwe_out.size(),
303                ct_grlwe_in.size(),
304                ct_rgsw.size(),
305            )
306            | GGSWCiphertext::encrypt_sk_scratch_space(&module, ct_rgsw.size()),
307    );
308
309    let k: usize = 1;
310
311    pt_rgsw.raw_mut()[k] = 1; // X^{k}
312
313    pt_grlwe.fill_ternary_prob(0, 0.5, &mut source_xs);
314
315    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
316    sk.fill_ternary_prob(0.5, &mut source_xs);
317
318    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
319    sk_dft.dft(&module, &sk);
320
321    // GRLWE_{s1}(s0) = s0 -> s1
322    ct_grlwe_in.encrypt_sk(
323        &module,
324        &pt_grlwe,
325        &sk_dft,
326        &mut source_xa,
327        &mut source_xe,
328        sigma,
329        bound,
330        scratch.borrow(),
331    );
332
333    ct_rgsw.encrypt_sk(
334        &module,
335        &pt_rgsw,
336        &sk_dft,
337        &mut source_xa,
338        &mut source_xe,
339        sigma,
340        bound,
341        scratch.borrow(),
342    );
343
344    // GRLWE_(m) (x) RGSW_(X^k) = GRLWE_(m * X^k)
345    ct_grlwe_out.external_product(&module, &ct_grlwe_in, &ct_rgsw, scratch.borrow());
346
347    let mut ct_rlwe_dft_s0s2: GLWECiphertextFourier<Vec<u8>, FFT64> =
348        GLWECiphertextFourier::new(&module, log_base2k, log_k_grlwe);
349    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_grlwe);
350
351    module.vec_znx_rotate_inplace(k as i64, &mut pt_grlwe, 0);
352
353    (0..ct_grlwe_out.rows()).for_each(|row_i| {
354        ct_grlwe_out.get_row(&module, row_i, 0, &mut ct_rlwe_dft_s0s2);
355        ct_rlwe_dft_s0s2.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
356        module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &pt_grlwe, 0);
357
358        let noise_have: f64 = pt.data.std(0, log_base2k).log2();
359
360        let var_gct_err_lhs: f64 = sigma * sigma;
361        let var_gct_err_rhs: f64 = 0f64;
362
363        let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
364        let var_a0_err: f64 = sigma * sigma;
365        let var_a1_err: f64 = 1f64 / 12f64;
366
367        let noise_want: f64 = noise_rgsw_product(
368            module.n() as f64,
369            log_base2k,
370            0.5,
371            var_msg,
372            var_a0_err,
373            var_a1_err,
374            var_gct_err_lhs,
375            var_gct_err_rhs,
376            log_k_grlwe,
377            log_k_grlwe,
378        );
379
380        assert!(
381            (noise_have - noise_want).abs() <= 0.1,
382            "{} {}",
383            noise_have,
384            noise_want
385        );
386    });
387}
388
389#[test]
390fn external_product_inplace() {
391    let module: Module<FFT64> = Module::<FFT64>::new(2048);
392    let log_base2k: usize = 12;
393    let log_k_grlwe: usize = 60;
394    let rows: usize = (log_k_grlwe + log_base2k - 1) / log_base2k;
395
396    let sigma: f64 = 3.2;
397    let bound: f64 = sigma * 6.0;
398
399    let mut ct_grlwe: GLWEKeySwitchKey<Vec<u8>, FFT64> = GLWEKeySwitchKey::new(&module, log_base2k, log_k_grlwe, rows);
400    let mut ct_rgsw: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::new(&module, log_base2k, log_k_grlwe, rows);
401
402    let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
403    let mut pt_grlwe: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
404
405    let mut source_xs: Source = Source::new([0u8; 32]);
406    let mut source_xe: Source = Source::new([0u8; 32]);
407    let mut source_xa: Source = Source::new([0u8; 32]);
408
409    let mut scratch: ScratchOwned = ScratchOwned::new(
410        GLWEKeySwitchKey::encrypt_sk_scratch_space(&module, ct_grlwe.size())
411            | GLWECiphertextFourier::decrypt_scratch_space(&module, ct_grlwe.size())
412            | GLWEKeySwitchKey::external_product_inplace_scratch_space(&module, ct_grlwe.size(), ct_rgsw.size())
413            | GGSWCiphertext::encrypt_sk_scratch_space(&module, ct_rgsw.size()),
414    );
415
416    let k: usize = 1;
417
418    pt_rgsw.raw_mut()[k] = 1; // X^{k}
419
420    pt_grlwe.fill_ternary_prob(0, 0.5, &mut source_xs);
421
422    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
423    sk.fill_ternary_prob(0.5, &mut source_xs);
424
425    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module);
426    sk_dft.dft(&module, &sk);
427
428    // GRLWE_{s1}(s0) = s0 -> s1
429    ct_grlwe.encrypt_sk(
430        &module,
431        &pt_grlwe,
432        &sk_dft,
433        &mut source_xa,
434        &mut source_xe,
435        sigma,
436        bound,
437        scratch.borrow(),
438    );
439
440    ct_rgsw.encrypt_sk(
441        &module,
442        &pt_rgsw,
443        &sk_dft,
444        &mut source_xa,
445        &mut source_xe,
446        sigma,
447        bound,
448        scratch.borrow(),
449    );
450
451    // GRLWE_(m) (x) RGSW_(X^k) = GRLWE_(m * X^k)
452    ct_grlwe.external_product_inplace(&module, &ct_rgsw, scratch.borrow());
453
454    let mut ct_rlwe_dft_s0s2: GLWECiphertextFourier<Vec<u8>, FFT64> =
455        GLWECiphertextFourier::new(&module, log_base2k, log_k_grlwe);
456    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, log_base2k, log_k_grlwe);
457
458    module.vec_znx_rotate_inplace(k as i64, &mut pt_grlwe, 0);
459
460    (0..ct_grlwe.rows()).for_each(|row_i| {
461        ct_grlwe.get_row(&module, row_i, 0, &mut ct_rlwe_dft_s0s2);
462        ct_rlwe_dft_s0s2.decrypt(&module, &mut pt, &sk_dft, scratch.borrow());
463        module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &pt_grlwe, 0);
464
465        let noise_have: f64 = pt.data.std(0, log_base2k).log2();
466
467        let var_gct_err_lhs: f64 = sigma * sigma;
468        let var_gct_err_rhs: f64 = 0f64;
469
470        let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
471        let var_a0_err: f64 = sigma * sigma;
472        let var_a1_err: f64 = 1f64 / 12f64;
473
474        let noise_want: f64 = noise_rgsw_product(
475            module.n() as f64,
476            log_base2k,
477            0.5,
478            var_msg,
479            var_a0_err,
480            var_a1_err,
481            var_gct_err_lhs,
482            var_gct_err_rhs,
483            log_k_grlwe,
484            log_k_grlwe,
485        );
486
487        assert!(
488            (noise_have - noise_want).abs() <= 0.1,
489            "{} {}",
490            noise_have,
491            noise_want
492        );
493    });
494}
495
496pub(crate) fn noise_grlwe_rlwe_product(
497    n: f64,
498    log_base2k: usize,
499    var_xs: f64,
500    var_msg: f64,
501    var_a_err: f64,
502    var_gct_err_lhs: f64,
503    var_gct_err_rhs: f64,
504    a_logq: usize,
505    b_logq: usize,
506) -> f64 {
507    let a_logq: usize = a_logq.min(b_logq);
508    let a_cols: usize = (a_logq + log_base2k - 1) / log_base2k;
509
510    let b_scale = 2.0f64.powi(b_logq as i32);
511    let a_scale: f64 = 2.0f64.powi((b_logq - a_logq) as i32);
512
513    let base: f64 = (1 << (log_base2k)) as f64;
514    let var_base: f64 = base * base / 12f64;
515
516    // lhs = a_cols * n * (var_base * var_gct_err_lhs + var_e_a * var_msg * p^2)
517    // rhs = a_cols * n * var_base * var_gct_err_rhs * var_xs
518    let mut noise: f64 = (a_cols as f64) * n * var_base * (var_gct_err_lhs + var_xs * var_gct_err_rhs);
519    noise += var_msg * var_a_err * a_scale * a_scale * n;
520    noise = noise.sqrt();
521    noise /= b_scale;
522    noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
523}