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