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