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