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