1use std::collections::HashMap;
  2
  3use backend::{FFT64, FillUniform, Module, ScratchOwned, Stats, VecZnxOps, ZnxView, ZnxViewMut};
  4use sampling::source::Source;
  5
  6use crate::{AutomorphismKey, GLWECiphertext, GLWEPlaintext, GLWESecret, Infos, test_fft64::gglwe::var_noise_gglwe_product};
  7
  8#[test]
  9fn trace_inplace() {
 10    (1..4).for_each(|rank| {
 11        println!("test trace_inplace rank: {}", rank);
 12        test_trace_inplace(11, 8, 54, 3.2, rank);
 13    });
 14}
 15
 16fn test_trace_inplace(log_n: usize, basek: usize, k: usize, sigma: f64, rank: usize) {
 17    let module: Module<FFT64> = Module::<FFT64>::new(1 << log_n);
 18
 19    let k_autokey: usize = k + basek;
 20
 21    let rows: usize = (k + basek - 1) / basek;
 22
 23    let mut ct: GLWECiphertext<Vec<u8>> = GLWECiphertext::alloc(&module, basek, k, rank);
 24    let mut pt_want: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
 25    let mut pt_have: GLWEPlaintext<Vec<u8>> = GLWEPlaintext::alloc(&module, basek, k);
 26
 27    let mut source_xs: Source = Source::new([0u8; 32]);
 28    let mut source_xe: Source = Source::new([0u8; 32]);
 29    let mut source_xa: Source = Source::new([0u8; 32]);
 30
 31    let mut scratch: ScratchOwned = ScratchOwned::new(
 32        GLWECiphertext::encrypt_sk_scratch_space(&module, basek, ct.k())
 33            | GLWECiphertext::decrypt_scratch_space(&module, basek, ct.k())
 34            | AutomorphismKey::generate_from_sk_scratch_space(&module, basek, k_autokey, rank)
 35            | GLWECiphertext::trace_inplace_scratch_space(&module, basek, ct.k(), k_autokey, rank),
 36    );
 37
 38    let mut sk: GLWESecret<Vec<u8>, FFT64> = GLWESecret::alloc(&module, rank);
 39    sk.fill_ternary_prob(&&module, 0.5, &mut source_xs);
 40
 41    let mut data_want: Vec<i64> = vec![0i64; module.n()];
 42
 43    data_want
 44        .iter_mut()
 45        .for_each(|x| *x = source_xa.next_i64() & 0xFF);
 46
 47    pt_have
 48        .data
 49        .fill_uniform(basek, 0, pt_have.size(), &mut source_xa);
 50
 51    ct.encrypt_sk(
 52        &module,
 53        &pt_have,
 54        &sk,
 55        &mut source_xa,
 56        &mut source_xe,
 57        sigma,
 58        scratch.borrow(),
 59    );
 60
 61    let mut auto_keys: HashMap<i64, AutomorphismKey<Vec<u8>, FFT64>> = HashMap::new();
 62    let gal_els: Vec<i64> = GLWECiphertext::trace_galois_elements(&module);
 63    gal_els.iter().for_each(|gal_el| {
 64        let mut key: AutomorphismKey<Vec<u8>, FFT64> = AutomorphismKey::alloc(&module, basek, k_autokey, rows, rank);
 65        key.generate_from_sk(
 66            &module,
 67            *gal_el,
 68            &sk,
 69            &mut source_xa,
 70            &mut source_xe,
 71            sigma,
 72            scratch.borrow(),
 73        );
 74        auto_keys.insert(*gal_el, key);
 75    });
 76
 77    ct.trace_inplace(&module, 0, 5, &auto_keys, scratch.borrow());
 78    ct.trace_inplace(&module, 5, log_n, &auto_keys, scratch.borrow());
 79
 80    (0..pt_want.size()).for_each(|i| pt_want.data.at_mut(0, i)[0] = pt_have.data.at(0, i)[0]);
 81
 82    ct.decrypt(&module, &mut pt_have, &sk, scratch.borrow());
 83
 84    module.vec_znx_sub_ab_inplace(&mut pt_want.data, 0, &pt_have.data, 0);
 85    module.vec_znx_normalize_inplace(basek, &mut pt_want.data, 0, scratch.borrow());
 86
 87    let noise_have = pt_want.data.std(0, basek).log2();
 88
 89    let mut noise_want: f64 = var_noise_gglwe_product(
 90        module.n() as f64,
 91        basek,
 92        0.5,
 93        0.5,
 94        1.0 / 12.0,
 95        sigma * sigma,
 96        0.0,
 97        rank as f64,
 98        k,
 99        k_autokey,
100    );
101    noise_want += sigma * sigma * (-2.0 * (k) as f64).exp2();
102    noise_want += module.n() as f64 * 1.0 / 12.0 * 0.5 * rank as f64 * (-2.0 * (k) as f64).exp2();
103    noise_want = noise_want.sqrt().log2();
104
105    assert!((noise_have - noise_want).abs() < 1.0);
106}