1use std::collections::HashMap;
 2
 3use backend::{FFT64, Module, Scratch};
 4
 5use crate::{AutomorphismKey, GLWECiphertext, GLWECiphertextToMut, GLWECiphertextToRef, GLWEOps, Infos, SetMetaData};
 6
 7impl GLWECiphertext<Vec<u8>> {
 8    pub fn trace_galois_elements(module: &Module<FFT64>) -> 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
20    pub fn trace_scratch_space(
21        module: &Module<FFT64>,
22        basek: usize,
23        out_k: usize,
24        in_k: usize,
25        atk_k: usize,
26        rank: usize,
27    ) -> usize {
28        Self::automorphism_inplace_scratch_space(module, basek, out_k.min(in_k), atk_k, rank)
29    }
30
31    pub fn trace_inplace_scratch_space(module: &Module<FFT64>, basek: usize, out_k: usize, atk_k: usize, rank: usize) -> usize {
32        Self::automorphism_inplace_scratch_space(module, basek, out_k, atk_k, rank)
33    }
34}
35
36impl<DataSelf: AsRef<[u8]> + AsMut<[u8]>> GLWECiphertext<DataSelf>
37where
38    GLWECiphertext<DataSelf>: GLWECiphertextToMut + Infos + SetMetaData,
39{
40    pub fn trace<DataLhs: AsRef<[u8]>, DataAK: AsRef<[u8]>>(
41        &mut self,
42        module: &Module<FFT64>,
43        start: usize,
44        end: usize,
45        lhs: &GLWECiphertext<DataLhs>,
46        auto_keys: &HashMap<i64, AutomorphismKey<DataAK, FFT64>>,
47        scratch: &mut Scratch,
48    ) where
49        GLWECiphertext<DataLhs>: GLWECiphertextToRef + Infos,
50    {
51        self.copy(module, lhs);
52        self.trace_inplace(module, start, end, auto_keys, scratch);
53    }
54
55    pub fn trace_inplace<DataAK: AsRef<[u8]>>(
56        &mut self,
57        module: &Module<FFT64>,
58        start: usize,
59        end: usize,
60        auto_keys: &HashMap<i64, AutomorphismKey<DataAK, FFT64>>,
61        scratch: &mut Scratch,
62    ) {
63        (start..end).for_each(|i| {
64            self.rsh(1, scratch);
65
66            let p: i64;
67            if i == 0 {
68                p = -1;
69            } else {
70                p = module.galois_element(1 << (i - 1));
71            }
72
73            if let Some(key) = auto_keys.get(&p) {
74                self.automorphism_add_inplace(module, key, scratch);
75            } else {
76                panic!("auto_keys[{}] is empty", p)
77            }
78        });
79    }
80}