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