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