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