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