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