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