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