1use base2k::{FFT64, Module, ScalarZnxOps, ScalarZnxToRef, ScratchOwned, Stats, VecZnxOps, ZnxView};
  2use sampling::source::Source;
  3
  4use crate::{
  5    automorphism::AutomorphismKey,
  6    elem::{GetRow, Infos},
  7    glwe_ciphertext_fourier::GLWECiphertextFourier,
  8    glwe_plaintext::GLWEPlaintext,
  9    keys::{SecretKey, SecretKeyFourier},
 10    test_fft64::gglwe::noise_gglwe_product,
 11};
 12
 13#[test]
 14fn automorphism() {
 15    (1..4).for_each(|rank| {
 16        println!("test automorphism rank: {}", rank);
 17        test_automorphism(-1, 5, 12, 12, 60, 3.2, rank);
 18    });
 19}
 20
 21fn test_automorphism(p0: i64, p1: i64, log_n: usize, basek: usize, k_ksk: usize, sigma: f64, rank: usize) {
 22    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 23    let rows = (k_ksk + basek - 1) / basek;
 24
 25    let mut auto_key_in: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, k_ksk, rows, rank);
 26    let mut auto_key_out: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, k_ksk, rows, rank);
 27    let mut auto_key_apply: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::new(&module, basek, k_ksk, rows, rank);
 28
 29    let mut source_xs: Source = Source::new([0u8; 32]);
 30    let mut source_xe: Source = Source::new([0u8; 32]);
 31    let mut source_xa: Source = Source::new([0u8; 32]);
 32
 33    let mut scratch: ScratchOwned = ScratchOwned::new(
 34        AutomorphismKey::encrypt_sk_scratch_space(&module, rank, auto_key_in.size())
 35            | GLWECiphertextFourier::decrypt_scratch_space(&module, auto_key_out.size())
 36            | AutomorphismKey::automorphism_scratch_space(
 37                &module,
 38                auto_key_out.size(),
 39                auto_key_in.size(),
 40                auto_key_apply.size(),
 41                rank,
 42            ),
 43    );
 44
 45    let mut sk: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
 46    sk.fill_ternary_prob(0.5, &mut source_xs);
 47
 48    let mut sk_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
 49    sk_dft.dft(&module, &sk);
 50
 51    // gglwe_{s1}(s0) = s0 -> s1
 52    auto_key_in.encrypt_sk(
 53        &module,
 54        p0,
 55        &sk,
 56        &mut source_xa,
 57        &mut source_xe,
 58        sigma,
 59        scratch.borrow(),
 60    );
 61
 62    // gglwe_{s2}(s1) -> s1 -> s2
 63    auto_key_apply.encrypt_sk(
 64        &module,
 65        p1,
 66        &sk,
 67        &mut source_xa,
 68        &mut source_xe,
 69        sigma,
 70        scratch.borrow(),
 71    );
 72
 73    // gglwe_{s1}(s0) (x) gglwe_{s2}(s1) = gglwe_{s2}(s0)
 74    auto_key_out.automorphism(&module, &auto_key_in, &auto_key_apply, scratch.borrow());
 75
 76    let mut ct_glwe_dft: GLWECiphertextFourier<Vec<u8>, FFT64> = GLWECiphertextFourier::new(&module, basek, k_ksk, rank);
 77    let mut pt: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::new(&module, basek, k_ksk);
 78
 79    let mut sk_auto: SecretKey<Vec<u8>> = SecretKey::new(&module, rank);
 80    sk_auto.fill_zero(); // Necessary to avoid panic of unfilled sk
 81    (0..rank).for_each(|i| {
 82        module.scalar_znx_automorphism(module.galois_element_inv(p0 * p1), &mut sk_auto, i, &sk, i);
 83    });
 84
 85    let mut sk_auto_dft: SecretKeyFourier<Vec<u8>, FFT64> = SecretKeyFourier::new(&module, rank);
 86    sk_auto_dft.dft(&module, &sk_auto);
 87
 88    (0..auto_key_out.rank_in()).for_each(|col_i| {
 89        (0..auto_key_out.rows()).for_each(|row_i| {
 90            auto_key_out.get_row(&module, row_i, col_i, &mut ct_glwe_dft);
 91
 92            ct_glwe_dft.decrypt(&module, &mut pt, &sk_auto_dft, scratch.borrow());
 93            module.vec_znx_sub_scalar_inplace(&mut pt, 0, row_i, &sk, col_i);
 94
 95            let noise_have: f64 = pt.data.std(0, basek).log2();
 96            let noise_want: f64 = noise_gglwe_product(
 97                module.n() as f64,
 98                basek,
 99                0.5,
100                0.5,
101                0f64,
102                sigma * sigma,
103                0f64,
104                rank as f64,
105                k_ksk,
106                k_ksk,
107            );
108
109            assert!(
110                (noise_have - noise_want).abs() <= 0.1,
111                "{} {}",
112                noise_have,
113                noise_want
114            );
115        });
116    });
117}