1use std::collections::HashMap;
 2
 3use backend::{FFT64, MatZnxDft, MatZnxDftToRef, Module, Scratch, VecZnx, VecZnxToMut, VecZnxToRef};
 4
 5use crate::{automorphism::AutomorphismKey, glwe_ciphertext::GLWECiphertext};
 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        out_size: usize,
23        in_size: usize,
24        autokey_size: usize,
25        rank: usize,
26    ) -> usize {
27        Self::automorphism_inplace_scratch_space(module, out_size.max(in_size), rank, autokey_size)
28    }
29
30    pub fn trace_inplace_scratch_space(module: &Module<FFT64>, out_size: usize, autokey_size: usize, rank: usize) -> usize {
31        Self::automorphism_inplace_scratch_space(module, out_size, rank, autokey_size)
32    }
33}
34
35impl<DataSelf> GLWECiphertext<DataSelf>
36where
37    VecZnx<DataSelf>: VecZnxToMut + VecZnxToRef,
38{
39    pub fn trace<DataLhs, DataAK>(
40        &mut self,
41        module: &Module<FFT64>,
42        start: usize,
43        end: usize,
44        lhs: &GLWECiphertext<DataLhs>,
45        auto_keys: &HashMap<i64, AutomorphismKey<DataAK, FFT64>>,
46        scratch: &mut Scratch,
47    ) where
48        VecZnx<DataLhs>: VecZnxToRef,
49        MatZnxDft<DataAK, FFT64>: MatZnxDftToRef<FFT64>,
50    {
51        self.copy(lhs);
52        self.trace_inplace(module, start, end, auto_keys, scratch);
53    }
54
55    pub fn trace_inplace<DataAK>(
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    ) where
63        MatZnxDft<DataAK, FFT64>: MatZnxDftToRef<FFT64>,
64    {
65        (start..end).for_each(|i| {
66            self.rsh(1, scratch);
67
68            let p: i64;
69            if i == 0 {
70                p = -1;
71            } else {
72                p = module.galois_element(1 << (i - 1));
73            }
74
75            if let Some(key) = auto_keys.get(&p) {
76                self.automorphism_add_inplace(module, key, scratch);
77            } else {
78                panic!("auto_keys[{}] is empty", p)
79            }
80        });
81    }
82}