1use crate::{
  2    elem::{Infos, ProdInplace, ProdInplaceScratchSpace, ProdScratchSpace, Product},
  3    grlwe::GRLWECt,
  4    keys::{SecretKey, SecretKeyDft},
  5    rgsw::RGSWCt,
  6    rlwe::{RLWECt, RLWECtDft, RLWEPt},
  7    test_fft64::{grlwe::noise_grlwe_rlwe_product, rgsw::noise_rgsw_product},
  8};
  9use base2k::{FFT64, FillUniform, Module, ScalarZnx, ScalarZnxAlloc, ScratchOwned, Stats, VecZnxOps, VecZnxToMut, ZnxViewMut};
 10use sampling::source::Source;
 11
 12#[test]
 13fn by_grlwe_inplace() {
 14    let module: Module<FFT64> = Module::<FFT64>::new(2048);
 15    let log_base2k: usize = 12;
 16    let log_k_grlwe: usize = 60;
 17    let log_k_rlwe_in: usize = 45;
 18    let log_k_rlwe_out: usize = 60;
 19    let rows: usize = (log_k_rlwe_in + log_base2k - 1) / log_base2k;
 20
 21    let sigma: f64 = 3.2;
 22    let bound: f64 = sigma * 6.0;
 23
 24    let mut ct_grlwe: GRLWECt<Vec<u8>, FFT64> = GRLWECt::new(&module, log_base2k, log_k_grlwe, rows);
 25    let mut ct_rlwe_in: RLWECt<Vec<u8>> = RLWECt::new(&module, log_base2k, log_k_rlwe_in);
 26    let mut ct_rlwe_in_dft: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_rlwe_in);
 27    let mut ct_rlwe_out: RLWECt<Vec<u8>> = RLWECt::new(&module, log_base2k, log_k_rlwe_out);
 28    let mut ct_rlwe_out_dft: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_rlwe_out);
 29    let mut pt_want: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_rlwe_in);
 30    let mut pt_have: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_rlwe_out);
 31
 32    let mut source_xs: Source = Source::new([0u8; 32]);
 33    let mut source_xe: Source = Source::new([0u8; 32]);
 34    let mut source_xa: Source = Source::new([0u8; 32]);
 35
 36    // Random input plaintext
 37    pt_want
 38        .data
 39        .fill_uniform(log_base2k, 0, pt_want.size(), &mut source_xa);
 40
 41    let mut scratch: ScratchOwned = ScratchOwned::new(
 42        GRLWECt::encrypt_sk_scratch_space(&module, ct_grlwe.size())
 43            | RLWECt::decrypt_scratch_space(&module, ct_rlwe_out.size())
 44            | RLWECt::encrypt_sk_scratch_space(&module, ct_rlwe_in.size())
 45            | RLWECtDft::prod_by_grlwe_scratch_space(
 46                &module,
 47                ct_rlwe_out.size(),
 48                ct_rlwe_in.size(),
 49                ct_grlwe.size(),
 50            ),
 51    );
 52
 53    let mut sk0: SecretKey<Vec<u8>> = SecretKey::new(&module);
 54    sk0.fill_ternary_prob(0.5, &mut source_xs);
 55
 56    let mut sk0_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
 57    sk0_dft.dft(&module, &sk0);
 58
 59    let mut sk1: SecretKey<Vec<u8>> = SecretKey::new(&module);
 60    sk1.fill_ternary_prob(0.5, &mut source_xs);
 61
 62    let mut sk1_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
 63    sk1_dft.dft(&module, &sk1);
 64
 65    ct_grlwe.encrypt_sk(
 66        &module,
 67        &sk0.data,
 68        &sk1_dft,
 69        &mut source_xa,
 70        &mut source_xe,
 71        sigma,
 72        bound,
 73        scratch.borrow(),
 74    );
 75
 76    ct_rlwe_in.encrypt_sk(
 77        &module,
 78        Some(&pt_want),
 79        &sk0_dft,
 80        &mut source_xa,
 81        &mut source_xe,
 82        sigma,
 83        bound,
 84        scratch.borrow(),
 85    );
 86
 87    ct_rlwe_in.dft(&module, &mut ct_rlwe_in_dft);
 88    ct_rlwe_out_dft.prod_by_grlwe(&module, &ct_rlwe_in_dft, &ct_grlwe, scratch.borrow());
 89    ct_rlwe_out_dft.idft(&module, &mut ct_rlwe_out, scratch.borrow());
 90
 91    ct_rlwe_out.decrypt(&module, &mut pt_have, &sk1_dft, scratch.borrow());
 92
 93    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
 94
 95    let noise_have: f64 = pt_have.data.std(0, log_base2k).log2();
 96    let noise_want: f64 = noise_grlwe_rlwe_product(
 97        module.n() as f64,
 98        log_base2k,
 99        0.5,
100        0.5,
101        0f64,
102        sigma * sigma,
103        0f64,
104        log_k_rlwe_in,
105        log_k_grlwe,
106    );
107
108    assert!(
109        (noise_have - noise_want).abs() <= 0.1,
110        "{} {}",
111        noise_have,
112        noise_want
113    );
114
115    module.free();
116}
117
118#[test]
119fn prod_by_grlwe_inplace() {
120    let module: Module<FFT64> = Module::<FFT64>::new(2048);
121    let log_base2k: usize = 12;
122    let log_k_grlwe: usize = 60;
123    let log_k_rlwe: usize = 45;
124    let rows: usize = (log_k_rlwe + log_base2k - 1) / log_base2k;
125
126    let sigma: f64 = 3.2;
127    let bound: f64 = sigma * 6.0;
128
129    let mut ct_grlwe: GRLWECt<Vec<u8>, FFT64> = GRLWECt::new(&module, log_base2k, log_k_grlwe, rows);
130    let mut ct_rlwe: RLWECt<Vec<u8>> = RLWECt::new(&module, log_base2k, log_k_rlwe);
131    let mut ct_rlwe_dft: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_rlwe);
132    let mut pt_want: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_rlwe);
133    let mut pt_have: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_rlwe);
134
135    let mut source_xs: Source = Source::new([0u8; 32]);
136    let mut source_xe: Source = Source::new([0u8; 32]);
137    let mut source_xa: Source = Source::new([0u8; 32]);
138
139    // Random input plaintext
140    pt_want
141        .data
142        .fill_uniform(log_base2k, 0, pt_want.size(), &mut source_xa);
143
144    let mut scratch: ScratchOwned = ScratchOwned::new(
145        GRLWECt::encrypt_sk_scratch_space(&module, ct_grlwe.size())
146            | RLWECt::decrypt_scratch_space(&module, ct_rlwe.size())
147            | RLWECt::encrypt_sk_scratch_space(&module, ct_rlwe.size())
148            | RLWECtDft::prod_by_grlwe_inplace_scratch_space(&module, ct_rlwe_dft.size(), ct_grlwe.size()),
149    );
150
151    let mut sk0: SecretKey<Vec<u8>> = SecretKey::new(&module);
152    sk0.fill_ternary_prob(0.5, &mut source_xs);
153
154    let mut sk0_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
155    sk0_dft.dft(&module, &sk0);
156
157    let mut sk1: SecretKey<Vec<u8>> = SecretKey::new(&module);
158    sk1.fill_ternary_prob(0.5, &mut source_xs);
159
160    let mut sk1_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
161    sk1_dft.dft(&module, &sk1);
162
163    ct_grlwe.encrypt_sk(
164        &module,
165        &sk0.data,
166        &sk1_dft,
167        &mut source_xa,
168        &mut source_xe,
169        sigma,
170        bound,
171        scratch.borrow(),
172    );
173
174    ct_rlwe.encrypt_sk(
175        &module,
176        Some(&pt_want),
177        &sk0_dft,
178        &mut source_xa,
179        &mut source_xe,
180        sigma,
181        bound,
182        scratch.borrow(),
183    );
184
185    ct_rlwe.dft(&module, &mut ct_rlwe_dft);
186    ct_rlwe_dft.prod_by_grlwe_inplace(&module, &ct_grlwe, scratch.borrow());
187    ct_rlwe_dft.idft(&module, &mut ct_rlwe, scratch.borrow());
188
189    ct_rlwe.decrypt(&module, &mut pt_have, &sk1_dft, scratch.borrow());
190
191    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
192
193    let noise_have: f64 = pt_have.data.std(0, log_base2k).log2();
194    let noise_want: f64 = noise_grlwe_rlwe_product(
195        module.n() as f64,
196        log_base2k,
197        0.5,
198        0.5,
199        0f64,
200        sigma * sigma,
201        0f64,
202        log_k_rlwe,
203        log_k_grlwe,
204    );
205
206    assert!(
207        (noise_have - noise_want).abs() <= 0.1,
208        "{} {}",
209        noise_have,
210        noise_want
211    );
212
213    module.free();
214}
215
216#[test]
217fn prod_by_rgsw() {
218    let module: Module<FFT64> = Module::<FFT64>::new(2048);
219    let log_base2k: usize = 12;
220    let log_k_grlwe: usize = 60;
221    let log_k_rlwe_in: usize = 45;
222    let log_k_rlwe_out: usize = 60;
223    let rows: usize = (log_k_rlwe_in + log_base2k - 1) / log_base2k;
224
225    let sigma: f64 = 3.2;
226    let bound: f64 = sigma * 6.0;
227
228    let mut ct_rgsw: RGSWCt<Vec<u8>, FFT64> = RGSWCt::new(&module, log_base2k, log_k_grlwe, rows);
229    let mut ct_rlwe_in: RLWECt<Vec<u8>> = RLWECt::new(&module, log_base2k, log_k_rlwe_in);
230    let mut ct_rlwe_out: RLWECt<Vec<u8>> = RLWECt::new(&module, log_base2k, log_k_rlwe_out);
231    let mut ct_rlwe_dft_in: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_rlwe_in);
232    let mut ct_rlwe_dft_out: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_rlwe_out);
233    let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
234    let mut pt_want: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_rlwe_in);
235    let mut pt_have: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_rlwe_out);
236
237    let mut source_xs: Source = Source::new([0u8; 32]);
238    let mut source_xe: Source = Source::new([0u8; 32]);
239    let mut source_xa: Source = Source::new([0u8; 32]);
240
241    // Random input plaintext
242    pt_want
243        .data
244        .fill_uniform(log_base2k, 0, pt_want.size(), &mut source_xa);
245
246    pt_want.to_mut().at_mut(0, 0)[1] = 1;
247
248    let k: usize = 1;
249
250    pt_rgsw.raw_mut()[k] = 1; // X^{k}
251
252    let mut scratch: ScratchOwned = ScratchOwned::new(
253        RGSWCt::encrypt_sk_scratch_space(&module, ct_rgsw.size())
254            | RLWECt::decrypt_scratch_space(&module, ct_rlwe_out.size())
255            | RLWECt::encrypt_sk_scratch_space(&module, ct_rlwe_in.size())
256            | RLWECt::prod_by_rgsw_scratch_space(
257                &module,
258                ct_rlwe_out.size(),
259                ct_rlwe_in.size(),
260                ct_rgsw.size(),
261            ),
262    );
263
264    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
265    sk.fill_ternary_prob(0.5, &mut source_xs);
266
267    let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
268    sk_dft.dft(&module, &sk);
269
270    ct_rgsw.encrypt_sk(
271        &module,
272        &pt_rgsw,
273        &sk_dft,
274        &mut source_xa,
275        &mut source_xe,
276        sigma,
277        bound,
278        scratch.borrow(),
279    );
280
281    ct_rlwe_in.encrypt_sk(
282        &module,
283        Some(&pt_want),
284        &sk_dft,
285        &mut source_xa,
286        &mut source_xe,
287        sigma,
288        bound,
289        scratch.borrow(),
290    );
291
292    ct_rlwe_in.dft(&module, &mut ct_rlwe_dft_in);
293    ct_rlwe_dft_out.prod_by_rgsw(&module, &ct_rlwe_dft_in, &ct_rgsw, scratch.borrow());
294    ct_rlwe_dft_out.idft(&module, &mut ct_rlwe_out, scratch.borrow());
295
296    ct_rlwe_out.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
297
298    module.vec_znx_rotate_inplace(k as i64, &mut pt_want, 0);
299
300    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
301
302    let noise_have: f64 = pt_have.data.std(0, log_base2k).log2();
303
304    let var_gct_err_lhs: f64 = sigma * sigma;
305    let var_gct_err_rhs: f64 = 0f64;
306
307    let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
308    let var_a0_err: f64 = sigma * sigma;
309    let var_a1_err: f64 = 1f64 / 12f64;
310
311    let noise_want: f64 = noise_rgsw_product(
312        module.n() as f64,
313        log_base2k,
314        0.5,
315        var_msg,
316        var_a0_err,
317        var_a1_err,
318        var_gct_err_lhs,
319        var_gct_err_rhs,
320        log_k_rlwe_in,
321        log_k_grlwe,
322    );
323
324    assert!(
325        (noise_have - noise_want).abs() <= 0.1,
326        "{} {}",
327        noise_have,
328        noise_want
329    );
330
331    module.free();
332}
333
334#[test]
335fn prod_by_rgsw_inplace() {
336    let module: Module<FFT64> = Module::<FFT64>::new(2048);
337    let log_base2k: usize = 12;
338    let log_k_grlwe: usize = 60;
339    let log_k_rlwe_in: usize = 45;
340    let log_k_rlwe_out: usize = 60;
341    let rows: usize = (log_k_rlwe_in + log_base2k - 1) / log_base2k;
342
343    let sigma: f64 = 3.2;
344    let bound: f64 = sigma * 6.0;
345
346    let mut ct_rgsw: RGSWCt<Vec<u8>, FFT64> = RGSWCt::new(&module, log_base2k, log_k_grlwe, rows);
347    let mut ct_rlwe: RLWECt<Vec<u8>> = RLWECt::new(&module, log_base2k, log_k_rlwe_in);
348    let mut ct_rlwe_dft: RLWECtDft<Vec<u8>, FFT64> = RLWECtDft::new(&module, log_base2k, log_k_rlwe_in);
349    let mut pt_rgsw: ScalarZnx<Vec<u8>> = module.new_scalar_znx(1);
350    let mut pt_want: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_rlwe_in);
351    let mut pt_have: RLWEPt<Vec<u8>> = RLWEPt::new(&module, log_base2k, log_k_rlwe_out);
352
353    let mut source_xs: Source = Source::new([0u8; 32]);
354    let mut source_xe: Source = Source::new([0u8; 32]);
355    let mut source_xa: Source = Source::new([0u8; 32]);
356
357    // Random input plaintext
358    pt_want
359        .data
360        .fill_uniform(log_base2k, 0, pt_want.size(), &mut source_xa);
361
362    pt_want.to_mut().at_mut(0, 0)[1] = 1;
363
364    let k: usize = 1;
365
366    pt_rgsw.raw_mut()[k] = 1; // X^{k}
367
368    let mut scratch: ScratchOwned = ScratchOwned::new(
369        RGSWCt::encrypt_sk_scratch_space(&module, ct_rgsw.size())
370            | RLWECt::decrypt_scratch_space(&module, ct_rlwe.size())
371            | RLWECt::encrypt_sk_scratch_space(&module, ct_rlwe.size())
372            | RLWECt::prod_by_rgsw_inplace_scratch_space(&module, ct_rlwe.size(), ct_rgsw.size()),
373    );
374
375    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module);
376    sk.fill_ternary_prob(0.5, &mut source_xs);
377
378    let mut sk_dft: SecretKeyDft<Vec<u8>, FFT64> = SecretKeyDft::new(&module);
379    sk_dft.dft(&module, &sk);
380
381    ct_rgsw.encrypt_sk(
382        &module,
383        &pt_rgsw,
384        &sk_dft,
385        &mut source_xa,
386        &mut source_xe,
387        sigma,
388        bound,
389        scratch.borrow(),
390    );
391
392    ct_rlwe.encrypt_sk(
393        &module,
394        Some(&pt_want),
395        &sk_dft,
396        &mut source_xa,
397        &mut source_xe,
398        sigma,
399        bound,
400        scratch.borrow(),
401    );
402
403    ct_rlwe.dft(&module, &mut ct_rlwe_dft);
404    ct_rlwe_dft.prod_by_rgsw_inplace(&module, &ct_rgsw, scratch.borrow());
405    ct_rlwe_dft.idft(&module, &mut ct_rlwe, scratch.borrow());
406
407    ct_rlwe.decrypt(&module, &mut pt_have, &sk_dft, scratch.borrow());
408
409    module.vec_znx_rotate_inplace(k as i64, &mut pt_want, 0);
410
411    module.vec_znx_sub_ab_inplace(&mut pt_have, 0, &pt_want, 0);
412
413    let noise_have: f64 = pt_have.data.std(0, log_base2k).log2();
414
415    let var_gct_err_lhs: f64 = sigma * sigma;
416    let var_gct_err_rhs: f64 = 0f64;
417
418    let var_msg: f64 = 1f64 / module.n() as f64; // X^{k}
419    let var_a0_err: f64 = sigma * sigma;
420    let var_a1_err: f64 = 1f64 / 12f64;
421
422    let noise_want: f64 = noise_rgsw_product(
423        module.n() as f64,
424        log_base2k,
425        0.5,
426        var_msg,
427        var_a0_err,
428        var_a1_err,
429        var_gct_err_lhs,
430        var_gct_err_rhs,
431        log_k_rlwe_in,
432        log_k_grlwe,
433    );
434
435    assert!(
436        (noise_have - noise_want).abs() <= 0.1,
437        "{} {}",
438        noise_have,
439        noise_want
440    );
441
442    module.free();
443}