1use std::collections::HashMap;
  2
  3use backend::hal::layouts::{Backend, DataMut, DataRef, Module, Scratch};
  4
  5use crate::{
  6    layouts::{GLWECiphertext, prepared::GGLWEAutomorphismKeyPrepared},
  7    operations::GLWEOperations,
  8};
  9
 10use crate::trait_families::{GLWETraceModuleFamily, GLWETraceScratchFamily};
 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>: GLWETraceModuleFamily<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>: GLWETraceModuleFamily<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> {
 58    pub fn trace<DataLhs: DataRef, DataAK: DataRef, B: Backend>(
 59        &mut self,
 60        module: &Module<B>,
 61        start: usize,
 62        end: usize,
 63        lhs: &GLWECiphertext<DataLhs>,
 64        auto_keys: &HashMap<i64, GGLWEAutomorphismKeyPrepared<DataAK, B>>,
 65        scratch: &mut Scratch<B>,
 66    ) where
 67        Module<B>: GLWETraceModuleFamily<B>,
 68        Scratch<B>: GLWETraceScratchFamily<B>,
 69    {
 70        self.copy(module, lhs);
 71        self.trace_inplace(module, start, end, auto_keys, scratch);
 72    }
 73
 74    pub fn trace_inplace<DataAK: DataRef, B: Backend>(
 75        &mut self,
 76        module: &Module<B>,
 77        start: usize,
 78        end: usize,
 79        auto_keys: &HashMap<i64, GGLWEAutomorphismKeyPrepared<DataAK, B>>,
 80        scratch: &mut Scratch<B>,
 81    ) where
 82        Module<B>: GLWETraceModuleFamily<B>,
 83        Scratch<B>: GLWETraceScratchFamily<B>,
 84    {
 85        (start..end).for_each(|i| {
 86            self.rsh(module, 1);
 87
 88            let p: i64;
 89            if i == 0 {
 90                p = -1;
 91            } else {
 92                p = module.galois_element(1 << (i - 1));
 93            }
 94
 95            if let Some(key) = auto_keys.get(&p) {
 96                self.automorphism_add_inplace(module, key, scratch);
 97            } else {
 98                panic!("auto_keys[{}] is empty", p)
 99            }
100        });
101    }
102}