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