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