1use std::collections::HashMap;
  2
  3use poulpy_hal::{
  4    api::{
  5        ScratchAvailable, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace, VecZnxBigNormalize, VecZnxBigNormalizeTmpBytes,
  6        VecZnxCopy, VecZnxDftApply, VecZnxDftBytesOf, VecZnxIdftApplyConsume, VecZnxNormalize, VecZnxNormalizeTmpBytes,
  7        VecZnxRshInplace, VmpApplyDftToDft, VmpApplyDftToDftAdd, VmpApplyDftToDftTmpBytes,
  8    },
  9    layouts::{Backend, DataMut, DataRef, Module, Scratch, VecZnx},
 10};
 11
 12use crate::{
 13    layouts::{Base2K, GGLWEInfos, GLWE, GLWEInfos, GLWELayout, LWEInfos, prepared::AutomorphismKeyPrepared},
 14    operations::GLWEOperations,
 15};
 16
 17impl GLWE<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_tmp_bytes<B: Backend, OUT, IN, KEY>(module: &Module<B>, out_infos: &OUT, in_infos: &IN, key_infos: &KEY) -> usize
 31    where
 32        OUT: GLWEInfos,
 33        IN: GLWEInfos,
 34        KEY: GGLWEInfos,
 35        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
 36    {
 37        let trace: usize = Self::automorphism_inplace_tmp_bytes(module, out_infos, key_infos);
 38        if in_infos.base2k() != key_infos.base2k() {
 39            let glwe_conv: usize = VecZnx::bytes_of(
 40                module.n(),
 41                (key_infos.rank_out() + 1).into(),
 42                out_infos.k().min(in_infos.k()).div_ceil(key_infos.base2k()) as usize,
 43            ) + module.vec_znx_normalize_tmp_bytes();
 44            return glwe_conv + trace;
 45        }
 46
 47        trace
 48    }
 49
 50    pub fn trace_inplace_tmp_bytes<B: Backend, OUT, KEY>(module: &Module<B>, out_infos: &OUT, key_infos: &KEY) -> usize
 51    where
 52        OUT: GLWEInfos,
 53        KEY: GGLWEInfos,
 54        Module<B>: VecZnxDftBytesOf + VmpApplyDftToDftTmpBytes + VecZnxBigNormalizeTmpBytes + VecZnxNormalizeTmpBytes,
 55    {
 56        Self::trace_tmp_bytes(module, out_infos, out_infos, key_infos)
 57    }
 58}
 59
 60impl<DataSelf: DataMut> GLWE<DataSelf> {
 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: &GLWE<DataLhs>,
 67        auto_keys: &HashMap<i64, AutomorphismKeyPrepared<DataAK, B>>,
 68        scratch: &mut Scratch<B>,
 69    ) where
 70        Module<B>: VecZnxDftBytesOf
 71            + VmpApplyDftToDftTmpBytes
 72            + VecZnxBigNormalizeTmpBytes
 73            + VmpApplyDftToDft<B>
 74            + VmpApplyDftToDftAdd<B>
 75            + VecZnxDftApply<B>
 76            + VecZnxIdftApplyConsume<B>
 77            + VecZnxBigAddSmallInplace<B>
 78            + VecZnxBigNormalize<B>
 79            + VecZnxBigAutomorphismInplace<B>
 80            + VecZnxRshInplace<B>
 81            + VecZnxCopy
 82            + VecZnxNormalizeTmpBytes
 83            + VecZnxNormalize<B>,
 84        Scratch<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, AutomorphismKeyPrepared<DataAK, B>>,
 96        scratch: &mut Scratch<B>,
 97    ) where
 98        Module<B>: VecZnxDftBytesOf
 99            + VmpApplyDftToDftTmpBytes
100            + VecZnxBigNormalizeTmpBytes
101            + VmpApplyDftToDft<B>
102            + VmpApplyDftToDftAdd<B>
103            + VecZnxDftApply<B>
104            + VecZnxIdftApplyConsume<B>
105            + VecZnxBigAddSmallInplace<B>
106            + VecZnxBigNormalize<B>
107            + VecZnxBigAutomorphismInplace<B>
108            + VecZnxRshInplace<B>
109            + VecZnxNormalizeTmpBytes
110            + VecZnxNormalize<B>,
111        Scratch<B>: ScratchAvailable,
112    {
113        let basek_ksk: Base2K = auto_keys
114            .get(auto_keys.keys().next().unwrap())
115            .unwrap()
116            .base2k();
117
118        #[cfg(debug_assertions)]
119        {
120            assert_eq!(self.n(), module.n() as u32);
121            assert!(start < end);
122            assert!(end <= module.log_n());
123            for key in auto_keys.values() {
124                assert_eq!(key.n(), module.n() as u32);
125                assert_eq!(key.base2k(), basek_ksk);
126                assert_eq!(key.rank_in(), self.rank());
127                assert_eq!(key.rank_out(), self.rank());
128            }
129        }
130
131        if self.base2k() != basek_ksk {
132            let (mut self_conv, scratch_1) = scratch.take_glwe_ct(&GLWELayout {
133                n: module.n().into(),
134                base2k: basek_ksk,
135                k: self.k(),
136                rank: self.rank(),
137            });
138
139            for j in 0..(self.rank() + 1).into() {
140                module.vec_znx_normalize(
141                    basek_ksk.into(),
142                    &mut self_conv.data,
143                    j,
144                    basek_ksk.into(),
145                    &self.data,
146                    j,
147                    scratch_1,
148                );
149            }
150
151            for i in start..end {
152                self_conv.rsh(module, 1, scratch_1);
153
154                let p: i64 = if i == 0 {
155                    -1
156                } else {
157                    module.galois_element(1 << (i - 1))
158                };
159
160                if let Some(key) = auto_keys.get(&p) {
161                    self_conv.automorphism_add_inplace(module, key, scratch_1);
162                } else {
163                    panic!("auto_keys[{p}] is empty")
164                }
165            }
166
167            for j in 0..(self.rank() + 1).into() {
168                module.vec_znx_normalize(
169                    self.base2k().into(),
170                    &mut self.data,
171                    j,
172                    basek_ksk.into(),
173                    &self_conv.data,
174                    j,
175                    scratch_1,
176                );
177            }
178        } else {
179            for i in start..end {
180                self.rsh(module, 1, scratch);
181
182                let p: i64 = if i == 0 {
183                    -1
184                } else {
185                    module.galois_element(1 << (i - 1))
186                };
187
188                if let Some(key) = auto_keys.get(&p) {
189                    self.automorphism_add_inplace(module, key, scratch);
190                } else {
191                    panic!("auto_keys[{p}] is empty")
192                }
193            }
194        }
195    }
196}