arnaucube
arnaucube/poulpyPublic Self-hosted repository
poulpy/core/src/test_fft64/ggsw.rs
core/src/test_fft64/ggsw.rsRaw
  1use backend::{
  2    FFT64, Module, ScalarZnx, ScalarZnxAlloc, ScalarZnxDftOps, ScalarZnxOps, ScratchOwned, Stats, VecZnxBig, VecZnxBigAlloc,
  3    VecZnxBigOps, VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxOps, VecZnxToMut, ZnxViewMut, ZnxZero,
  4};
  5use sampling::source::Source;
  6
  7use crate::{
  8    GGSWCiphertext, GLWECiphertextFourier, GLWEPlaintext, GLWESecret, GLWESwitchingKey, GetRow, Infos, TensorKey,
  9    automorphism::AutomorphismKey,
 10};
 11
 12use super::gglwe::var_noise_gglwe_product;
 13
 14#[test]
 15fn encrypt_sk() {
 16    (1..4).for_each(|rank| {
 17        println!("test encrypt_sk rank: {}", rank);
 18        test_encrypt_sk(11, 8, 54, 3.2, rank);
 19    });
 20}
 21
 22#[test]
 23fn keyswitch() {
 24    (1..4).for_each(|rank| {
 25        println!("test keyswitch rank: {}", rank);
 26        test_keyswitch(12, 15, 60, rank, 3.2);
 27    });
 28}
 29
 30#[test]
 31fn keyswitch_inplace() {
 32    (1..4).for_each(|rank| {
 33        println!("test keyswitch_inplace rank: {}", rank);
 34        test_keyswitch_inplace(12, 15, 60, rank, 3.2);
 35    });
 36}
 37
 38#[test]
 39fn automorphism() {
 40    (1..4).for_each(|rank| {
 41        println!("test automorphism rank: {}", rank);
 42        test_automorphism(-5, 12, 15, 60, rank, 3.2);
 43    });
 44}
 45
 46#[test]
 47fn automorphism_inplace() {
 48    (1..4).for_each(|rank| {
 49        println!("test automorphism_inplace rank: {}", rank);
 50        test_automorphism_inplace(-5, 12, 15, 60, rank, 3.2);
 51    });
 52}
 53
 54#[test]
 55fn external_product() {
 56    (1..4).for_each(|rank| {
 57        println!("test external_product rank: {}", rank);
 58        test_external_product(12, 12, 60, rank, 3.2);
 59    });
 60}
 61
 62#[test]
 63fn external_product_inplace() {
 64    (1..4).for_each(|rank| {
 65        println!("test external_product rank: {}", rank);
 66        test_external_product_inplace(12, 15, 60, rank, 3.2);
 67    });
 68}
 69
 70fn test_encrypt_sk(log_n: usize, basek: usize, k: usize, sigma: f64, rank: usize) {
 71    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 72
 73    let rows: usize = (k + basek - 1) / basek;
 74
 75    let mut ct: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k, rows, rank);
 76    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
 77    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
 78    let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
 79
 80    let mut source_xs: Source = Source::new([0u8; 32]);
 81    let mut source_xe: Source = Source::new([0u8; 32]);
 82    let mut source_xa: Source = Source::new([0u8; 32]);
 83
 84    pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
 85
 86    let mut scratch: ScratchOwned = ScratchOwned::new(
 87        GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k, rank)
 88            | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k),
 89    );
 90
 91    let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
 92    sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
 93
 94    ct.encrypt_sk(
 95        &module,
 96        &pt_scalar,
 97        &sk,
 98        &mut source_xa,
 99        &mut source_xe,
100        sigma,
101        scratch.borrow(),
102    );
103
104    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k, rank);
105    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
106    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
107
108    (0..ct.rank() + 1).for_each(|col_j| {
109        (0..ct.rows()).for_each(|row_i| {
110            module.vec_znx_add_scalar_inplace(&mut pt_want.data, 0, row_i, &pt_scalar, 0);
111
112            // mul with sk[col_j-1]
113            if col_j > 0 {
114                module.vec_znx_dft(&mut pt_dft, 0, &pt_want.data, 0);
115                module.svp_apply_inplace(&mut pt_dft, 0, &sk.data_fourier, col_j - 1);
116                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
117                module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
118            }
119
120            ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
121
122            ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk, scratch.borrow());
123
124            module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
125
126            let std_pt: f64 = pt_have.data.std(0, basek) * (k as f64).exp2();
127            assert!((sigma - std_pt).abs() <= 0.2, "{} {}", sigma, std_pt);
128
129            pt_want.data.zero();
130        });
131    });
132}
133
134fn test_keyswitch(log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) {
135    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
136    let rows: usize = (k + basek - 1) / basek;
137
138    let mut ct_in: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k, rows, rank);
139    let mut ct_out: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k, rows, rank);
140    let mut tsk: TensorKey<Vec<u8>, FFT64> = TensorKey::alloc(&module, basek, k, rows, rank);
141    let mut ksk: GLWESwitchingKey<Vec<u8>, FFT64> = GLWESwitchingKey::alloc(&module, basek, k, rows, rank, rank);
142    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
143    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
144    let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
145
146    let mut source_xs: Source = Source::new([0u8; 32]);
147    let mut source_xe: Source = Source::new([0u8; 32]);
148    let mut source_xa: Source = Source::new([0u8; 32]);
149
150    let mut scratch: ScratchOwned = ScratchOwned::new(
151        GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k, rank)
152            | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k)
153            | GLWESwitchingKey::encrypt_sk_scratch_space(&module, basek, k, rank)
154            | TensorKey::generate_from_sk_scratch_space(&module, basek, k, rank)
155            | GGSWCiphertext::keyswitch_scratch_space(
156                &module,
157                basek,
158                ct_out.k(),
159                ct_in.k(),
160                ksk.k(),
161                tsk.k(),
162                rank,
163            ),
164    );
165
166    let var_xs: f64 = 0.5;
167
168    let mut sk_in: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
169    sk_in.fill_ternary_prob(&module, var_xs, &mut source_xs);
170
171    let mut sk_out: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
172    sk_out.fill_ternary_prob(&module, var_xs, &mut source_xs);
173
174    ksk.generate_from_sk(
175        &module,
176        &sk_in,
177        &sk_out,
178        &mut source_xa,
179        &mut source_xe,
180        sigma,
181        scratch.borrow(),
182    );
183    tsk.generate_from_sk(
184        &module,
185        &sk_out,
186        &mut source_xa,
187        &mut source_xe,
188        sigma,
189        scratch.borrow(),
190    );
191
192    pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
193
194    ct_in.encrypt_sk(
195        &module,
196        &pt_scalar,
197        &sk_in,
198        &mut source_xa,
199        &mut source_xe,
200        sigma,
201        scratch.borrow(),
202    );
203
204    ct_out.keyswitch(&module, &ct_in, &ksk, &tsk, scratch.borrow());
205
206    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k, rank);
207    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_out.size());
208    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_out.size());
209
210    (0..ct_out.rank() + 1).for_each(|col_j| {
211        (0..ct_out.rows()).for_each(|row_i| {
212            module.vec_znx_add_scalar_inplace(&mut pt_want.data, 0, row_i, &pt_scalar, 0);
213
214            // mul with sk[col_j-1]
215            if col_j > 0 {
216                module.vec_znx_dft(&mut pt_dft, 0, &pt_want.data, 0);
217                module.svp_apply_inplace(&mut pt_dft, 0, &sk_out.data_fourier, col_j - 1);
218                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
219                module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
220            }
221
222            ct_out.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
223
224            ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk_out, scratch.borrow());
225
226            module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
227
228            let noise_have: f64 = pt_have.data.std(0, basek).log2();
229            let noise_want: f64 = noise_ggsw_keyswitch(
230                module.n() as f64,
231                basek,
232                col_j,
233                var_xs,
234                0f64,
235                sigma * sigma,
236                0f64,
237                rank as f64,
238                k,
239                k,
240            );
241
242            println!("{} {}", noise_have, noise_want);
243
244            assert!(
245                (noise_have - noise_want).abs() <= 0.1,
246                "{} {}",
247                noise_have,
248                noise_want
249            );
250
251            pt_want.data.zero();
252        });
253    });
254}
255
256fn test_keyswitch_inplace(log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) {
257    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
258    let rows: usize = (k + basek - 1) / basek;
259
260    let mut ct: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k, rows, rank);
261    let mut tsk: TensorKey<Vec<u8>, FFT64> = TensorKey::alloc(&module, basek, k, rows, rank);
262    let mut ksk: GLWESwitchingKey<Vec<u8>, FFT64> = GLWESwitchingKey::alloc(&module, basek, k, rows, rank, rank);
263    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
264    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
265    let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
266
267    let mut source_xs: Source = Source::new([0u8; 32]);
268    let mut source_xe: Source = Source::new([0u8; 32]);
269    let mut source_xa: Source = Source::new([0u8; 32]);
270
271    let mut scratch: ScratchOwned = ScratchOwned::new(
272        GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k, rank)
273            | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k)
274            | GLWESwitchingKey::encrypt_sk_scratch_space(&module, basek, k, rank)
275            | TensorKey::generate_from_sk_scratch_space(&module, basek, k, rank)
276            | GGSWCiphertext::keyswitch_inplace_scratch_space(&module, basek, ct.k(), ksk.k(), tsk.k(), rank),
277    );
278
279    let var_xs: f64 = 0.5;
280
281    let mut sk_in: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
282    sk_in.fill_ternary_prob(&module, var_xs, &mut source_xs);
283
284    let mut sk_out: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
285    sk_out.fill_ternary_prob(&module, var_xs, &mut source_xs);
286
287    ksk.generate_from_sk(
288        &module,
289        &sk_in,
290        &sk_out,
291        &mut source_xa,
292        &mut source_xe,
293        sigma,
294        scratch.borrow(),
295    );
296    tsk.generate_from_sk(
297        &module,
298        &sk_out,
299        &mut source_xa,
300        &mut source_xe,
301        sigma,
302        scratch.borrow(),
303    );
304
305    pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
306
307    ct.encrypt_sk(
308        &module,
309        &pt_scalar,
310        &sk_in,
311        &mut source_xa,
312        &mut source_xe,
313        sigma,
314        scratch.borrow(),
315    );
316
317    ct.keyswitch_inplace(&module, &ksk, &tsk, scratch.borrow());
318
319    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k, rank);
320    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
321    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
322
323    (0..ct.rank() + 1).for_each(|col_j| {
324        (0..ct.rows()).for_each(|row_i| {
325            module.vec_znx_add_scalar_inplace(&mut pt_want.data, 0, row_i, &pt_scalar, 0);
326
327            // mul with sk[col_j-1]
328            if col_j > 0 {
329                module.vec_znx_dft(&mut pt_dft, 0, &pt_want.data, 0);
330                module.svp_apply_inplace(&mut pt_dft, 0, &sk_out.data_fourier, col_j - 1);
331                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
332                module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
333            }
334
335            ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
336
337            ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk_out, scratch.borrow());
338
339            module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
340
341            let noise_have: f64 = pt_have.data.std(0, basek).log2();
342            let noise_want: f64 = noise_ggsw_keyswitch(
343                module.n() as f64,
344                basek,
345                col_j,
346                var_xs,
347                0f64,
348                sigma * sigma,
349                0f64,
350                rank as f64,
351                k,
352                k,
353            );
354
355            println!("{} {}", noise_have, noise_want);
356
357            assert!(
358                (noise_have - noise_want).abs() <= 0.1,
359                "{} {}",
360                noise_have,
361                noise_want
362            );
363
364            pt_want.data.zero();
365        });
366    });
367}
368
369pub(crate) fn noise_ggsw_keyswitch(
370    n: f64,
371    basek: usize,
372    col: usize,
373    var_xs: f64,
374    var_a_err: f64,
375    var_gct_err_lhs: f64,
376    var_gct_err_rhs: f64,
377    rank: f64,
378    a_logq: usize,
379    b_logq: usize,
380) -> f64 {
381    let var_si_x_sj: f64 = n * var_xs * var_xs;
382
383    // Initial KS for col = 0
384    let mut noise: f64 = var_noise_gglwe_product(
385        n,
386        basek,
387        var_xs,
388        var_xs,
389        var_a_err,
390        var_gct_err_lhs,
391        var_gct_err_rhs,
392        rank,
393        a_logq,
394        b_logq,
395    );
396
397    // Other GGSW reconstruction for col > 0
398    if col > 0 {
399        noise += var_noise_gglwe_product(
400            n,
401            basek,
402            var_xs,
403            var_si_x_sj,
404            var_a_err + 1f64 / 12.0,
405            var_gct_err_lhs,
406            var_gct_err_rhs,
407            rank,
408            a_logq,
409            b_logq,
410        );
411        noise += n * noise * var_xs * 0.5;
412    }
413
414    noise = noise.sqrt();
415    noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
416}
417
418fn test_automorphism(p: i64, log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) {
419    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
420    let rows: usize = (k + basek - 1) / basek;
421
422    let mut ct_in: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k, rows, rank);
423    let mut ct_out: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k, rows, rank);
424    let mut tensor_key: TensorKey<Vec<u8>, FFT64> = TensorKey::alloc(&module, basek, k, rows, rank);
425    let mut auto_key: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::alloc(&module, basek, k, rows, rank);
426    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
427    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
428    let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
429
430    let mut source_xs: Source = Source::new([0u8; 32]);
431    let mut source_xe: Source = Source::new([0u8; 32]);
432    let mut source_xa: Source = Source::new([0u8; 32]);
433
434    let mut scratch: ScratchOwned = ScratchOwned::new(
435        GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k, rank)
436            | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k)
437            | AutomorphismKey::generate_from_sk_scratch_space(&module, basek, k, rank)
438            | TensorKey::generate_from_sk_scratch_space(&module, basek, k, rank)
439            | GGSWCiphertext::automorphism_scratch_space(
440                &module,
441                basek,
442                ct_out.k(),
443                ct_in.k(),
444                auto_key.k(),
445                tensor_key.k(),
446                rank,
447            ),
448    );
449
450    let var_xs: f64 = 0.5;
451
452    let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
453    sk.fill_ternary_prob(&module, var_xs, &mut source_xs);
454
455    auto_key.generate_from_sk(
456        &module,
457        p,
458        &sk,
459        &mut source_xa,
460        &mut source_xe,
461        sigma,
462        scratch.borrow(),
463    );
464    tensor_key.generate_from_sk(
465        &module,
466        &sk,
467        &mut source_xa,
468        &mut source_xe,
469        sigma,
470        scratch.borrow(),
471    );
472
473    pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
474
475    ct_in.encrypt_sk(
476        &module,
477        &pt_scalar,
478        &sk,
479        &mut source_xa,
480        &mut source_xe,
481        sigma,
482        scratch.borrow(),
483    );
484
485    ct_out.automorphism(&module, &ct_in, &auto_key, &tensor_key, scratch.borrow());
486
487    module.scalar_znx_automorphism_inplace(p, &mut pt_scalar, 0);
488
489    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k, rank);
490    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_out.size());
491    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_out.size());
492
493    (0..ct_out.rank() + 1).for_each(|col_j| {
494        (0..ct_out.rows()).for_each(|row_i| {
495            module.vec_znx_add_scalar_inplace(&mut pt_want.data, 0, row_i, &pt_scalar, 0);
496
497            // mul with sk[col_j-1]
498            if col_j > 0 {
499                module.vec_znx_dft(&mut pt_dft, 0, &pt_want.data, 0);
500                module.svp_apply_inplace(&mut pt_dft, 0, &sk.data_fourier, col_j - 1);
501                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
502                module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
503            }
504
505            ct_out.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
506
507            ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk, scratch.borrow());
508
509            module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
510
511            let noise_have: f64 = pt_have.data.std(0, basek).log2();
512            let noise_want: f64 = noise_ggsw_keyswitch(
513                module.n() as f64,
514                basek,
515                col_j,
516                var_xs,
517                0f64,
518                sigma * sigma,
519                0f64,
520                rank as f64,
521                k,
522                k,
523            );
524
525            assert!(
526                (noise_have - noise_want).abs() <= 0.1,
527                "{} {}",
528                noise_have,
529                noise_want
530            );
531
532            pt_want.data.zero();
533        });
534    });
535}
536
537fn test_automorphism_inplace(p: i64, log_n: usize, basek: usize, k: usize, rank: usize, sigma: f64) {
538    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
539    let rows: usize = (k + basek - 1) / basek;
540
541    let mut ct: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k, rows, rank);
542    let mut tensor_key: TensorKey<Vec<u8>, FFT64> = TensorKey::alloc(&module, basek, k, rows, rank);
543    let mut auto_key: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::alloc(&module, basek, k, rows, rank);
544    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
545    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
546    let mut pt_scalar: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
547
548    let mut source_xs: Source = Source::new([0u8; 32]);
549    let mut source_xe: Source = Source::new([0u8; 32]);
550    let mut source_xa: Source = Source::new([0u8; 32]);
551
552    let mut scratch: ScratchOwned = ScratchOwned::new(
553        GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k, rank)
554            | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k)
555            | AutomorphismKey::generate_from_sk_scratch_space(&module, basek, k, rank)
556            | TensorKey::generate_from_sk_scratch_space(&module, basek, k, rank)
557            | GGSWCiphertext::automorphism_inplace_scratch_space(&module, basek, ct.k(), auto_key.k(), tensor_key.k(), rank),
558    );
559
560    let var_xs: f64 = 0.5;
561
562    let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
563    sk.fill_ternary_prob(&module, var_xs, &mut source_xs);
564
565    auto_key.generate_from_sk(
566        &module,
567        p,
568        &sk,
569        &mut source_xa,
570        &mut source_xe,
571        sigma,
572        scratch.borrow(),
573    );
574    tensor_key.generate_from_sk(
575        &module,
576        &sk,
577        &mut source_xa,
578        &mut source_xe,
579        sigma,
580        scratch.borrow(),
581    );
582
583    pt_scalar.fill_ternary_hw(0, module.n(), &mut source_xs);
584
585    ct.encrypt_sk(
586        &module,
587        &pt_scalar,
588        &sk,
589        &mut source_xa,
590        &mut source_xe,
591        sigma,
592        scratch.borrow(),
593    );
594
595    ct.automorphism_inplace(&module, &auto_key, &tensor_key, scratch.borrow());
596
597    module.scalar_znx_automorphism_inplace(p, &mut pt_scalar, 0);
598
599    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k, rank);
600    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct.size());
601    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct.size());
602
603    (0..ct.rank() + 1).for_each(|col_j| {
604        (0..ct.rows()).for_each(|row_i| {
605            module.vec_znx_add_scalar_inplace(&mut pt_want.data, 0, row_i, &pt_scalar, 0);
606
607            // mul with sk[col_j-1]
608            if col_j > 0 {
609                module.vec_znx_dft(&mut pt_dft, 0, &pt_want.data, 0);
610                module.svp_apply_inplace(&mut pt_dft, 0, &sk.data_fourier, col_j - 1);
611                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
612                module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
613            }
614
615            ct.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
616
617            ct_glwe_fourier.decrypt(&module, &mut pt_have, &sk, scratch.borrow());
618
619            module.vec_znx_sub_ab_inplace(&mut pt_have.data, 0, &pt_want.data, 0);
620
621            let noise_have: f64 = pt_have.data.std(0, basek).log2();
622            let noise_want: f64 = noise_ggsw_keyswitch(
623                module.n() as f64,
624                basek,
625                col_j,
626                var_xs,
627                0f64,
628                sigma * sigma,
629                0f64,
630                rank as f64,
631                k,
632                k,
633            );
634
635            assert!(
636                (noise_have - noise_want).abs() <= 0.1,
637                "{} {}",
638                noise_have,
639                noise_want
640            );
641
642            pt_want.data.zero();
643        });
644    });
645}
646
647fn test_external_product(log_n: usize, basek: usize, k_ggsw: usize, rank: usize, sigma: f64) {
648    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
649
650    let rows: usize = (k_ggsw + basek - 1) / basek;
651
652    let mut ct_ggsw_rhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, rank);
653    let mut ct_ggsw_lhs_in: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, rank);
654    let mut ct_ggsw_lhs_out: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, rank);
655    let mut pt_ggsw_lhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
656    let mut pt_ggsw_rhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
657
658    let mut source_xs: Source = Source::new([0u8; 32]);
659    let mut source_xe: Source = Source::new([0u8; 32]);
660    let mut source_xa: Source = Source::new([0u8; 32]);
661
662    pt_ggsw_lhs.fill_ternary_prob(0, 0.5, &mut source_xs);
663
664    let k: usize = 1;
665
666    pt_ggsw_rhs.to_mut().raw_mut()[k] = 1; //X^{k}
667
668    let mut scratch: ScratchOwned = ScratchOwned::new(
669        GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k)
670            | GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k, rank)
671            | GGSWCiphertext::external_product_scratch_space(
672                &module,
673                basek,
674                ct_ggsw_lhs_out.k(),
675                ct_ggsw_lhs_in.k(),
676                ct_ggsw_rhs.k(),
677                rank,
678            ),
679    );
680
681    let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
682    sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
683
684    ct_ggsw_rhs.encrypt_sk(
685        &module,
686        &pt_ggsw_rhs,
687        &sk,
688        &mut source_xa,
689        &mut source_xe,
690        sigma,
691        scratch.borrow(),
692    );
693
694    ct_ggsw_lhs_in.encrypt_sk(
695        &module,
696        &pt_ggsw_lhs,
697        &sk,
698        &mut source_xa,
699        &mut source_xe,
700        sigma,
701        scratch.borrow(),
702    );
703
704    ct_ggsw_lhs_out.external_product(&module, &ct_ggsw_lhs_in, &ct_ggsw_rhs, scratch.borrow());
705
706    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_ggsw, rank);
707    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ggsw);
708    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_ggsw_lhs_out.size());
709    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_ggsw_lhs_out.size());
710    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ggsw);
711
712    module.vec_znx_rotate_inplace(k as i64, &mut pt_ggsw_lhs, 0);
713
714    (0..ct_ggsw_lhs_out.rank() + 1).for_each(|col_j| {
715        (0..ct_ggsw_lhs_out.rows()).for_each(|row_i| {
716            module.vec_znx_add_scalar_inplace(&mut pt_want.data, 0, row_i, &pt_ggsw_lhs, 0);
717
718            if col_j > 0 {
719                module.vec_znx_dft(&mut pt_dft, 0, &pt_want.data, 0);
720                module.svp_apply_inplace(&mut pt_dft, 0, &sk.data_fourier, col_j - 1);
721                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
722                module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
723            }
724
725            ct_ggsw_lhs_out.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
726            ct_glwe_fourier.decrypt(&module, &mut pt, &sk, scratch.borrow());
727
728            module.vec_znx_sub_ab_inplace(&mut pt.data, 0, &pt_want.data, 0);
729
730            let noise_have: f64 = pt.data.std(0, basek).log2();
731
732            let var_gct_err_lhs: f64 = sigma * sigma;
733            let var_gct_err_rhs: f64 = 0f64;
734
735            let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
736            let var_a0_err: f64 = sigma * sigma;
737            let var_a1_err: f64 = 1f64 / 12f64;
738
739            let noise_want: f64 = noise_ggsw_product(
740                module.n() as f64,
741                basek,
742                0.5,
743                var_msg,
744                var_a0_err,
745                var_a1_err,
746                var_gct_err_lhs,
747                var_gct_err_rhs,
748                rank as f64,
749                k_ggsw,
750                k_ggsw,
751            );
752
753            assert!(
754                (noise_have - noise_want).abs() <= 0.1,
755                "have: {} want: {}",
756                noise_have,
757                noise_want
758            );
759
760            pt_want.data.zero();
761        });
762    });
763}
764
765fn test_external_product_inplace(log_n: usize, basek: usize, k_ggsw: usize, rank: usize, sigma: f64) {
766    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
767    let rows: usize = (k_ggsw + basek - 1) / basek;
768
769    let mut ct_ggsw_rhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, rank);
770    let mut ct_ggsw_lhs: GGSWCiphertext<Vec<u8>, FFT64> = GGSWCiphertext::alloc(&module, basek, k_ggsw, rows, rank);
771    let mut pt_ggsw_lhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
772    let mut pt_ggsw_rhs: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
773
774    let mut source_xs: Source = Source::new([0u8; 32]);
775    let mut source_xe: Source = Source::new([0u8; 32]);
776    let mut source_xa: Source = Source::new([0u8; 32]);
777
778    pt_ggsw_lhs.fill_ternary_prob(0, 0.5, &mut source_xs);
779
780    let k: usize = 1;
781
782    pt_ggsw_rhs.to_mut().raw_mut()[k] = 1; //X^{k}
783
784    let mut scratch: ScratchOwned = ScratchOwned::new(
785        GLWESwitchingKey::encrypt_sk_scratch_space(&module, basek, k, rank)
786            | GLWECiphertextFourier::decrypt_scratch_space(&module, basek, k)
787            | GGSWCiphertext::encrypt_sk_scratch_space(&module, basek, k, rank)
788            | GGSWCiphertext::external_product_inplace_scratch_space(&module, basek, ct_ggsw_lhs.k(), ct_ggsw_rhs.k(), rank),
789    );
790
791    let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
792    sk.fill_ternary_prob(&module, 0.5, &mut source_xs);
793
794    ct_ggsw_rhs.encrypt_sk(
795        &module,
796        &pt_ggsw_rhs,
797        &sk,
798        &mut source_xa,
799        &mut source_xe,
800        sigma,
801        scratch.borrow(),
802    );
803
804    ct_ggsw_lhs.encrypt_sk(
805        &module,
806        &pt_ggsw_lhs,
807        &sk,
808        &mut source_xa,
809        &mut source_xe,
810        sigma,
811        scratch.borrow(),
812    );
813
814    ct_ggsw_lhs.external_product_inplace(&module, &ct_ggsw_rhs, scratch.borrow());
815
816    let mut ct_glwe_fourier: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::alloc(&module, basek, k_ggsw, rank);
817    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ggsw);
818    let mut pt_dft: VecZnxDft<Vec<u8>, FFT64> = module.new_vec_znx_dft(1, ct_ggsw_lhs.size());
819    let mut pt_big: VecZnxBig<Vec<u8>, FFT64> = module.new_vec_znx_big(1, ct_ggsw_lhs.size());
820    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k_ggsw);
821
822    module.vec_znx_rotate_inplace(k as i64, &mut pt_ggsw_lhs, 0);
823
824    (0..ct_ggsw_lhs.rank() + 1).for_each(|col_j| {
825        (0..ct_ggsw_lhs.rows()).for_each(|row_i| {
826            module.vec_znx_add_scalar_inplace(&mut pt_want.data, 0, row_i, &pt_ggsw_lhs, 0);
827
828            if col_j > 0 {
829                module.vec_znx_dft(&mut pt_dft, 0, &pt_want.data, 0);
830                module.svp_apply_inplace(&mut pt_dft, 0, &sk.data_fourier, col_j - 1);
831                module.vec_znx_idft_tmp_a(&mut pt_big, 0, &mut pt_dft, 0);
832                module.vec_znx_big_normalize(basek, &mut pt_want.data, 0, &pt_big, 0, scratch.borrow());
833            }
834
835            ct_ggsw_lhs.get_row(&module, row_i, col_j, &mut ct_glwe_fourier);
836            ct_glwe_fourier.decrypt(&module, &mut pt, &sk, scratch.borrow());
837
838            module.vec_znx_sub_ab_inplace(&mut pt.data, 0, &pt_want.data, 0);
839
840            let noise_have: f64 = pt.data.std(0, basek).log2();
841
842            let var_gct_err_lhs: f64 = sigma * sigma;
843            let var_gct_err_rhs: f64 = 0f64;
844
845            let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
846            let var_a0_err: f64 = sigma * sigma;
847            let var_a1_err: f64 = 1f64 / 12f64;
848
849            let noise_want: f64 = noise_ggsw_product(
850                module.n() as f64,
851                basek,
852                0.5,
853                var_msg,
854                var_a0_err,
855                var_a1_err,
856                var_gct_err_lhs,
857                var_gct_err_rhs,
858                rank as f64,
859                k_ggsw,
860                k_ggsw,
861            );
862
863            assert!(
864                (noise_have - noise_want).abs() <= 0.1,
865                "have: {} want: {}",
866                noise_have,
867                noise_want
868            );
869
870            pt_want.data.zero();
871        });
872    });
873}
874pub(crate) fn noise_ggsw_product(
875    n: f64,
876    basek: usize,
877    var_xs: f64,
878    var_msg: f64,
879    var_a0_err: f64,
880    var_a1_err: f64,
881    var_gct_err_lhs: f64,
882    var_gct_err_rhs: f64,
883    rank: f64,
884    a_logq: usize,
885    b_logq: usize,
886) -> f64 {
887    let a_logq: usize = a_logq.min(b_logq);
888    let a_cols: usize = (a_logq + basek - 1) / basek;
889
890    let b_scale = 2.0f64.powi(b_logq as i32);
891    let a_scale: f64 = 2.0f64.powi((b_logq - a_logq) as i32);
892
893    let base: f64 = (1 << (basek)) as f64;
894    let var_base: f64 = base * base / 12f64;
895
896    // lhs = a_cols * n * (var_base * var_gct_err_lhs + var_e_a * var_msg * p^2)
897    // rhs = a_cols * n * var_base * var_gct_err_rhs * var_xs
898    let mut noise: f64 = (rank + 1.0) * (a_cols as f64) * n * var_base * (var_gct_err_lhs + var_xs * var_gct_err_rhs);
899    noise += var_msg * var_a0_err * a_scale * a_scale * n;
900    noise += var_msg * var_a1_err * a_scale * a_scale * n * var_xs * rank;
901    noise = noise.sqrt();
902    noise /= b_scale;
903    noise.log2().min(-1.0) // max noise is [-2^{-1}, 2^{-1}]
904}