1use crate::{automorphism::AutomorphismKey, ciphertext::Ciphertext, elem::ElemCommon};
  2use base2k::{
  3    Module, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VecZnxOps, VmpPMatOps,
  4    assert_alignement,
  5};
  6use std::collections::HashMap;
  7
  8pub fn trace_tmp_bytes(
  9    module: &Module,
 10    c_cols: usize,
 11    a_cols: usize,
 12    b_rows: usize,
 13    b_cols: usize,
 14) -> usize {
 15    return module.vmp_apply_dft_to_dft_tmp_bytes(c_cols, a_cols, b_rows, b_cols)
 16        + 2 * module.bytes_of_vec_znx_dft(std::cmp::min(c_cols, a_cols));
 17}
 18
 19pub fn trace_inplace(
 20    module: &Module,
 21    a: &mut Ciphertext<VecZnx>,
 22    start: usize,
 23    end: usize,
 24    b: HashMap<i64, AutomorphismKey>,
 25    tmp_bytes: &mut [u8],
 26) {
 27    let cols: usize = a.cols();
 28
 29    let b_rows: usize;
 30    let b_cols: usize;
 31
 32    if let Some((_, key)) = b.iter().next() {
 33        b_rows = key.value.rows();
 34        b_cols = key.value.cols();
 35    } else {
 36        panic!("b: HashMap<i64, AutomorphismKey>, is empty")
 37    }
 38
 39    #[cfg(debug_assertions)]
 40    {
 41        assert!(start <= end);
 42        assert!(end <= module.n());
 43        assert!(tmp_bytes.len() >= trace_tmp_bytes(module, cols, cols, b_rows, b_cols));
 44        assert_alignement(tmp_bytes.as_ptr());
 45    }
 46
 47    let cols: usize = std::cmp::min(b_cols, a.cols());
 48
 49    let (tmp_bytes_b1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
 50    let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
 51
 52    let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_b1_dft);
 53    let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_res_dft);
 54    let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
 55
 56    let log_base2k: usize = a.log_base2k();
 57
 58    (start..end).for_each(|i| {
 59        a.at_mut(0).rsh(log_base2k, 1, tmp_bytes);
 60        a.at_mut(1).rsh(log_base2k, 1, tmp_bytes);
 61
 62        let p: i64;
 63        if i == 0 {
 64            p = -1;
 65        } else {
 66            p = module.galois_element(1 << (i - 1));
 67        }
 68
 69        if let Some(key) = b.get(&p) {
 70            module.vec_znx_dft(&mut a1_dft, a.at(1));
 71
 72            // a[0] = NORMALIZE(a[0] + AUTO(a[0] + IDFT(<DFT(a[1]), key[0]>)))
 73            module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, key.value.at(0), tmp_bytes);
 74            module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
 75            module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0));
 76            module.vec_znx_big_automorphism_inplace(p, &mut res_big);
 77            module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0));
 78            module.vec_znx_big_normalize(a.log_base2k(), a.at_mut(0), &mut res_big, tmp_bytes);
 79
 80            // a[1] = NORMALIZE(a[1] + AUTO(IDFT(<DFT(a[1]), key[1]>)))
 81            module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, key.value.at(1), tmp_bytes);
 82            module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
 83            module.vec_znx_big_automorphism_inplace(p, &mut res_big);
 84            module.vec_znx_big_add_small_inplace(&mut res_big, a.at(1));
 85            module.vec_znx_big_normalize(a.log_base2k(), a.at_mut(1), &mut res_big, tmp_bytes);
 86        } else {
 87            panic!("b[{}] is empty", p)
 88        }
 89    })
 90}
 91
 92#[cfg(test)]
 93mod test {
 94    use crate::{
 95        ciphertext::{Ciphertext, new_gadget_ciphertext},
 96        decryptor::decrypt_rlwe,
 97        elem::{Elem, ElemCommon, ElemVecZnx},
 98        encryptor::encrypt_grlwe_sk,
 99        gadget_product::gadget_product_core,
100        keys::SecretKey,
101        parameters::{Parameters, ParametersLiteral},
102        plaintext::Plaintext,
103    };
104    use base2k::{
105        BACKEND, Infos, Sampling, SvpPPolOps, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft,
106        VecZnxDftOps, VecZnxOps, VmpPMat, alloc_aligned_u8,
107    };
108    use sampling::source::{Source, new_seed};
109
110    #[test]
111    fn test_trace_inplace() {}
112}