1use crate::{automorphism::AutomorphismKey, ciphertext::Ciphertext, elem::ElemCommon, parameters::Parameters};
  2use base2k::{
  3    Module, VecZnx, VecZnxBig, VecZnxBigOps, VecZnxDft, VecZnxDftOps, VmpPMatOps,
  4    assert_alignement,
  5};
  6use std::collections::HashMap;
  7
  8pub fn trace_galois_elements(module: &Module) -> Vec<i64>{
  9    let mut gal_els: Vec<i64> = Vec::new();
 10    (0..module.log_n()).for_each(|i|{
 11        if i == 0 {
 12            gal_els.push(-1);
 13        } else {
 14            gal_els.push(module.galois_element(1 << (i - 1)));
 15        }
 16    });
 17    gal_els
 18}
 19
 20impl Parameters{
 21    pub fn trace_tmp_bytes(&self, res_logq: usize, in_logq: usize, gct_logq: usize) -> usize{
 22        self.automorphism_tmp_bytes(res_logq, in_logq, gct_logq)
 23    }
 24}
 25
 26pub fn trace_tmp_bytes(
 27    module: &Module,
 28    c_cols: usize,
 29    a_cols: usize,
 30    b_rows: usize,
 31    b_cols: usize,
 32) -> usize {
 33    return module.vmp_apply_dft_to_dft_tmp_bytes(c_cols, a_cols, b_rows, b_cols)
 34        + 2 * module.bytes_of_vec_znx_dft(std::cmp::min(c_cols, a_cols));
 35}
 36
 37pub fn trace_inplace(
 38    module: &Module,
 39    a: &mut Ciphertext<VecZnx>,
 40    start: usize,
 41    end: usize,
 42    b: &HashMap<i64, AutomorphismKey>,
 43    b_cols: usize,
 44    tmp_bytes: &mut [u8],
 45) {
 46    let cols: usize = a.cols();
 47
 48    let b_rows: usize;
 49
 50    if let Some((_, key)) = b.iter().next() {
 51        b_rows = key.value.rows();
 52        #[cfg(debug_assertions)]{
 53            println!("{} {}", b_cols, key.value.cols());
 54            assert!(b_cols <= key.value.cols())
 55        }
 56    } else {
 57        panic!("b: HashMap<i64, AutomorphismKey>, is empty")
 58    }
 59
 60    #[cfg(debug_assertions)]
 61    {
 62        assert!(start <= end);
 63        assert!(end <= module.n());
 64        assert!(tmp_bytes.len() >= trace_tmp_bytes(module, cols, cols, b_rows, b_cols));
 65        assert_alignement(tmp_bytes.as_ptr());
 66    }
 67
 68    let cols: usize = std::cmp::min(b_cols, a.cols());
 69
 70    let (tmp_bytes_b1_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(cols));
 71    let (tmp_bytes_res_dft, tmp_bytes) = tmp_bytes.split_at_mut(module.bytes_of_vec_znx_dft(b_cols));
 72
 73    let mut a1_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(cols, tmp_bytes_b1_dft);
 74    let mut res_dft: VecZnxDft = module.new_vec_znx_dft_from_bytes_borrow(b_cols, tmp_bytes_res_dft);
 75    let mut res_big: VecZnxBig = res_dft.as_vec_znx_big();
 76
 77    let log_base2k: usize = a.log_base2k();
 78
 79    (start..end).for_each(|i| {
 80        a.at_mut(0).rsh(log_base2k, 1, tmp_bytes);
 81        a.at_mut(1).rsh(log_base2k, 1, tmp_bytes);
 82
 83        let p: i64;
 84        if i == 0 {
 85            p = -1;
 86        } else {
 87            p = module.galois_element(1 << (i - 1));
 88        }
 89
 90        if let Some(key) = b.get(&p) {
 91            module.vec_znx_dft(&mut a1_dft, a.at(1));
 92
 93            // a[0] = NORMALIZE(a[0] + AUTO(a[0] + IDFT(<DFT(a[1]), key[0]>)))
 94            module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, key.value.at(0), tmp_bytes);
 95            module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
 96            module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0));
 97            module.vec_znx_big_automorphism_inplace(p, &mut res_big);
 98            module.vec_znx_big_add_small_inplace(&mut res_big, a.at(0));
 99            module.vec_znx_big_normalize(a.log_base2k(), a.at_mut(0), &mut res_big, tmp_bytes);
100
101            // a[1] = NORMALIZE(a[1] + AUTO(IDFT(<DFT(a[1]), key[1]>)))
102            module.vmp_apply_dft_to_dft(&mut res_dft, &a1_dft, key.value.at(1), tmp_bytes);
103            module.vec_znx_idft_tmp_a(&mut res_big, &mut res_dft);
104            module.vec_znx_big_automorphism_inplace(p, &mut res_big);
105            module.vec_znx_big_add_small_inplace(&mut res_big, a.at(1));
106            module.vec_znx_big_normalize(a.log_base2k(), a.at_mut(1), &mut res_big, tmp_bytes);
107        } else {
108            panic!("b[{}] is empty", p)
109        }
110    })
111}
112
113#[cfg(test)]
114mod test {
115    use crate::{
116        automorphism::AutomorphismKey, ciphertext::{new_gadget_ciphertext, Ciphertext}, decryptor::decrypt_rlwe, elem::{Elem, ElemCommon, ElemVecZnx}, encryptor::{encrypt_grlwe_sk, encrypt_rlwe_sk}, gadget_product::gadget_product_core, keys::SecretKey, parameters::{Parameters, ParametersLiteral, DEFAULTSIGMA}, plaintext::Plaintext
117    };
118    use base2k::{
119        BACKEND, Module, Infos, Sampling, SvpPPol, SvpPPolOps, VecZnx,
120        VecZnxDftOps, VecZnxOps, VmpPMat, alloc_aligned, Encoding,
121    };
122    use sampling::source::{Source, new_seed};
123    use std::collections::HashMap;
124    use super::{trace_galois_elements, trace_inplace};
125
126    #[test]
127    fn test_trace_inplace() {
128        let log_base2k: usize = 10;
129        let log_q: usize = 50;
130        let log_p: usize = 15;
131
132        // Basic parameters with enough limbs to test edge cases
133        let params_lit: ParametersLiteral = ParametersLiteral {
134            backend: BACKEND::FFT64,
135            log_n: 12,
136            log_q: log_q,
137            log_p: log_p,
138            log_base2k: log_base2k,
139            log_scale: 20,
140            xe: 3.2,
141            xs: 1 << 11,
142        };
143
144        let params: Parameters = Parameters::new(&params_lit);
145
146        let module: &Module = params.module();
147        let log_q: usize = params.log_q();
148        let log_qp: usize = params.log_qp();
149        let gct_rows: usize = params.cols_q();
150        let gct_cols: usize = params.cols_qp();
151
152        // scratch space
153        let mut tmp_bytes: Vec<u8> = alloc_aligned(
154            params.decrypt_rlwe_tmp_byte(log_q)
155                | params.encrypt_rlwe_sk_tmp_bytes(log_q)
156                | params.automorphism_key_new_tmp_bytes(gct_rows, log_qp)
157                | params.automorphism_tmp_bytes(log_q, log_q, log_qp),
158        );
159
160        // Samplers for public and private randomness
161        let mut source_xe: Source = Source::new(new_seed());
162        let mut source_xa: Source = Source::new(new_seed());
163        let mut source_xs: Source = Source::new(new_seed());
164
165        let mut sk: SecretKey = SecretKey::new(module);
166        sk.fill_ternary_hw(params.xs(), &mut source_xs);
167        let mut sk_svp_ppol: SvpPPol = module.new_svp_ppol();
168        module.svp_prepare(&mut sk_svp_ppol, &sk.0);
169
170        let gal_els: Vec<i64> = trace_galois_elements(module);
171
172        let auto_keys: HashMap<i64, AutomorphismKey> = AutomorphismKey::new_many(module, &gal_els, &sk, log_base2k, gct_rows, log_qp, &mut source_xa, &mut source_xe, DEFAULTSIGMA, &mut tmp_bytes);
173
174        let mut data: Vec<i64> = vec![0i64; params.n()];
175
176        data.iter_mut().enumerate().for_each(|(i, x)| *x = 1+i as i64);
177
178        let log_k: usize = 2 * log_base2k;
179
180        let mut ct: Ciphertext<VecZnx> = params.new_ciphertext(log_q);
181        let mut pt: Plaintext = params.new_plaintext(log_q);
182
183        pt.at_mut(0).encode_vec_i64(log_base2k, log_k, &data, 32);
184        pt.at_mut(0).normalize(log_base2k, &mut tmp_bytes);
185
186        pt.at(0).decode_vec_i64(log_base2k, log_k, &mut data);
187
188        pt.at(0).print(pt.cols(), 16);
189
190        encrypt_rlwe_sk(
191            module,
192            &mut ct.elem_mut(),
193            Some(pt.elem()),
194            &sk_svp_ppol,
195            &mut source_xa,
196            &mut source_xe,
197            params.xe(),
198            &mut tmp_bytes,
199        );
200
201        trace_inplace(module, &mut ct, module.log_n()-2, module.log_n(), &auto_keys, gct_cols, & mut tmp_bytes);
202
203        // pt = dec(auto(ct)) - auto(pt)
204        decrypt_rlwe(
205            module,
206            pt.elem_mut(),
207            ct.elem(),
208            &sk_svp_ppol,
209            &mut tmp_bytes,
210        );
211
212        pt.at(0).print(pt.cols(), 16);
213
214        pt.at(0).decode_vec_i64(log_base2k, log_k, &mut data);
215
216        println!("trace: {:?}", &data[..16]);
217
218    }
219}