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