1use std::collections::HashMap;
  2
  3use backend::hal::{
  4    api::{ScratchAvailable, TakeVecZnxDft, VecZnxBigAutomorphismInplace, VecZnxCopy, VecZnxRshInplace},
  5    layouts::{Backend, DataMut, DataRef, Module, Scratch},
  6};
  7
  8use crate::{
  9    AutomorphismKeyExec, GLWECiphertext, GLWECiphertextToMut, GLWECiphertextToRef, GLWEKeyswitchFamily, GLWEOps, Infos,
 10    SetMetaData,
 11};
 12
 13pub trait GLWETraceFamily<B: Backend> = GLWEKeyswitchFamily<B> + VecZnxCopy + VecZnxRshInplace + VecZnxBigAutomorphismInplace<B>;
 14
 15impl GLWECiphertext<Vec<u8>> {
 16    pub fn trace_galois_elements<B: Backend>(module: &Module<B>) -> Vec<i64> {
 17        let mut gal_els: Vec<i64> = Vec::new();
 18        (0..module.log_n()).for_each(|i| {
 19            if i == 0 {
 20                gal_els.push(-1);
 21            } else {
 22                gal_els.push(module.galois_element(1 << (i - 1)));
 23            }
 24        });
 25        gal_els
 26    }
 27
 28    pub fn trace_scratch_space<B: Backend>(
 29        module: &Module<B>,
 30        n: usize,
 31        basek: usize,
 32        out_k: usize,
 33        in_k: usize,
 34        ksk_k: usize,
 35        digits: usize,
 36        rank: usize,
 37    ) -> usize
 38    where
 39        Module<B>: GLWEKeyswitchFamily<B>,
 40    {
 41        Self::automorphism_inplace_scratch_space(module, n, basek, out_k.min(in_k), ksk_k, digits, rank)
 42    }
 43
 44    pub fn trace_inplace_scratch_space<B: Backend>(
 45        module: &Module<B>,
 46        n: usize,
 47        basek: usize,
 48        out_k: usize,
 49        ksk_k: usize,
 50        digits: usize,
 51        rank: usize,
 52    ) -> usize
 53    where
 54        Module<B>: GLWEKeyswitchFamily<B>,
 55    {
 56        Self::automorphism_inplace_scratch_space(module, n, basek, out_k, ksk_k, digits, rank)
 57    }
 58}
 59
 60impl<DataSelf: DataMut> GLWECiphertext<DataSelf>
 61where
 62    GLWECiphertext<DataSelf>: GLWECiphertextToMut + Infos + SetMetaData,
 63{
 64    pub fn trace<DataLhs: DataRef, DataAK: DataRef, B: Backend>(
 65        &mut self,
 66        module: &Module<B>,
 67        start: usize,
 68        end: usize,
 69        lhs: &GLWECiphertext<DataLhs>,
 70        auto_keys: &HashMap<i64, AutomorphismKeyExec<DataAK, B>>,
 71        scratch: &mut Scratch<B>,
 72    ) where
 73        GLWECiphertext<DataLhs>: GLWECiphertextToRef + Infos + VecZnxRshInplace,
 74        Module<B>: GLWETraceFamily<B>,
 75        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
 76    {
 77        self.copy(module, lhs);
 78        self.trace_inplace(module, start, end, auto_keys, scratch);
 79    }
 80
 81    pub fn trace_inplace<DataAK: DataRef, B: Backend>(
 82        &mut self,
 83        module: &Module<B>,
 84        start: usize,
 85        end: usize,
 86        auto_keys: &HashMap<i64, AutomorphismKeyExec<DataAK, B>>,
 87        scratch: &mut Scratch<B>,
 88    ) where
 89        Module<B>: GLWETraceFamily<B>,
 90        Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable,
 91    {
 92        (start..end).for_each(|i| {
 93            self.rsh(module, 1);
 94
 95            let p: i64;
 96            if i == 0 {
 97                p = -1;
 98            } else {
 99                p = module.galois_element(1 << (i - 1));
100            }
101
102            if let Some(key) = auto_keys.get(&p) {
103                self.automorphism_add_inplace(module, key, scratch);
104            } else {
105                panic!("auto_keys[{}] is empty", p)
106            }
107        });
108    }
109}