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