1use poulpy_hal::{
  2    api::{ModuleLogN, VecZnxNormalize, VecZnxNormalizeTmpBytes},
  3    layouts::{Backend, CyclotomicOrder, DataMut, GaloisElement, Module, Scratch, VecZnx, galois_element},
  4};
  5
  6use crate::{
  7    GLWEAutomorphism, GLWECopy, GLWEShift, ScratchTakeCore,
  8    layouts::{
  9        GGLWEInfos, GGLWELayout, GGLWEPreparedToRef, GLWE, GLWEAutomorphismKeyHelper, GLWEInfos, GLWELayout, GLWEToMut,
 10        GLWEToRef, GetGaloisElement, LWEInfos,
 11    },
 12};
 13
 14impl GLWE<Vec<u8>> {
 15    pub fn trace_galois_elements<M, BE: Backend>(module: &M) -> Vec<i64>
 16    where
 17        M: GLWETrace<BE>,
 18    {
 19        module.glwe_trace_galois_elements()
 20    }
 21
 22    pub fn trace_tmp_bytes<R, A, K, M, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
 23    where
 24        R: GLWEInfos,
 25        A: GLWEInfos,
 26        K: GGLWEInfos,
 27        M: GLWETrace<BE>,
 28    {
 29        module.glwe_trace_tmp_bytes(res_infos, a_infos, key_infos)
 30    }
 31}
 32
 33impl<D: DataMut> GLWE<D> {
 34    pub fn trace<A, H, K, M, BE: Backend>(&mut self, module: &M, skip: usize, a: &A, keys: &H, scratch: &mut Scratch<BE>)
 35    where
 36        A: GLWEToRef + GLWEInfos,
 37        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
 38        H: GLWEAutomorphismKeyHelper<K, BE>,
 39        Scratch<BE>: ScratchTakeCore<BE>,
 40        M: GLWETrace<BE>,
 41    {
 42        module.glwe_trace(self, skip, a, keys, scratch);
 43    }
 44
 45    pub fn trace_inplace<H, K, M, BE: Backend>(&mut self, module: &M, skip: usize, keys: &H, scratch: &mut Scratch<BE>)
 46    where
 47        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
 48        H: GLWEAutomorphismKeyHelper<K, BE>,
 49        Scratch<BE>: ScratchTakeCore<BE>,
 50        M: GLWETrace<BE>,
 51    {
 52        module.glwe_trace_inplace(self, skip, keys, scratch);
 53    }
 54}
 55
 56#[inline(always)]
 57pub fn trace_galois_elements(log_n: usize, cyclotomic_order: i64) -> Vec<i64> {
 58    (0..log_n)
 59        .map(|i| {
 60            if i == 0 {
 61                -1
 62            } else {
 63                galois_element(1 << (i - 1), cyclotomic_order)
 64            }
 65        })
 66        .collect()
 67}
 68
 69impl<BE: Backend> GLWETrace<BE> for Module<BE>
 70where
 71    Self: ModuleLogN
 72        + GaloisElement
 73        + GLWEAutomorphism<BE>
 74        + GLWEShift<BE>
 75        + GLWECopy
 76        + CyclotomicOrder
 77        + VecZnxNormalizeTmpBytes
 78        + VecZnxNormalize<BE>,
 79    Scratch<BE>: ScratchTakeCore<BE>,
 80{
 81    fn glwe_trace_galois_elements(&self) -> Vec<i64> {
 82        trace_galois_elements(self.log_n(), self.cyclotomic_order())
 83    }
 84
 85    fn glwe_trace_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
 86    where
 87        R: GLWEInfos,
 88        A: GLWEInfos,
 89        K: GGLWEInfos,
 90    {
 91        let trace: usize = self.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos);
 92        if a_infos.base2k() != key_infos.base2k() {
 93            let glwe_conv: usize = VecZnx::bytes_of(
 94                self.n(),
 95                (key_infos.rank_out() + 1).into(),
 96                res_infos.k().min(a_infos.k()).div_ceil(key_infos.base2k()) as usize,
 97            ) + self.vec_znx_normalize_tmp_bytes();
 98            return glwe_conv + trace;
 99        }
100
101        let tmp = if res_infos.k() > a_infos.k() {
102            GLWE::bytes_of_from_infos(res_infos)
103        } else {
104            GLWE::bytes_of_from_infos(a_infos)
105        };
106
107        trace + tmp
108    }
109
110    fn glwe_trace<R, A, K, H>(&self, res: &mut R, skip: usize, a: &A, keys: &H, scratch: &mut Scratch<BE>)
111    where
112        R: GLWEToMut + GLWEInfos,
113        A: GLWEToRef + GLWEInfos,
114        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
115        H: GLWEAutomorphismKeyHelper<K, BE>,
116    {
117        let (mut tmp, scratch_1) = if a.k() > res.k() {
118            scratch.take_glwe(a)
119        } else {
120            scratch.take_glwe(res)
121        };
122
123        self.glwe_copy(&mut tmp, a);
124        self.glwe_trace_inplace(&mut tmp, skip, keys, scratch_1);
125        self.glwe_copy(res, &tmp);
126    }
127
128    fn glwe_trace_inplace<R, K, H>(&self, res: &mut R, skip: usize, keys: &H, scratch: &mut Scratch<BE>)
129    where
130        R: GLWEToMut,
131        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
132        H: GLWEAutomorphismKeyHelper<K, BE>,
133    {
134        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
135
136        let ksk_infos: &GGLWELayout = &keys.automorphism_key_infos();
137        let log_n: usize = self.log_n();
138
139        assert_eq!(res.n(), self.n() as u32);
140        assert_eq!(ksk_infos.n(), self.n() as u32);
141        assert!(skip <= log_n);
142        assert_eq!(ksk_infos.rank_in(), res.rank());
143        assert_eq!(ksk_infos.rank_out(), res.rank());
144
145        if res.base2k() != ksk_infos.base2k() {
146            let (mut self_conv, scratch_1) = scratch.take_glwe(&GLWELayout {
147                n: self.n().into(),
148                base2k: ksk_infos.base2k(),
149                k: res.k(),
150                rank: res.rank(),
151            });
152
153            for j in 0..(res.rank() + 1).into() {
154                self.vec_znx_normalize(
155                    ksk_infos.base2k().into(),
156                    &mut self_conv.data,
157                    j,
158                    res.base2k().into(),
159                    res.data(),
160                    j,
161                    scratch_1,
162                );
163            }
164
165            for i in skip..log_n {
166                self.glwe_rsh(1, &mut self_conv, scratch_1);
167
168                let p: i64 = if i == 0 {
169                    -1
170                } else {
171                    self.galois_element(1 << (i - 1))
172                };
173
174                if let Some(key) = keys.get_automorphism_key(p) {
175                    self.glwe_automorphism_add_inplace(&mut self_conv, key, scratch_1);
176                } else {
177                    panic!("keys[{p}] is empty")
178                }
179            }
180
181            for j in 0..(res.rank() + 1).into() {
182                self.vec_znx_normalize(
183                    res.base2k().into(),
184                    res.data_mut(),
185                    j,
186                    ksk_infos.base2k().into(),
187                    &self_conv.data,
188                    j,
189                    scratch_1,
190                );
191            }
192        } else {
193            for i in skip..log_n {
194                self.glwe_rsh(1, res, scratch);
195
196                let p: i64 = if i == 0 {
197                    -1
198                } else {
199                    self.galois_element(1 << (i - 1))
200                };
201
202                if let Some(key) = keys.get_automorphism_key(p) {
203                    self.glwe_automorphism_add_inplace(res, key, scratch);
204                } else {
205                    panic!("keys[{p}] is empty")
206                }
207            }
208        }
209    }
210}
211
212pub trait GLWETrace<BE: Backend> {
213    fn glwe_trace_galois_elements(&self) -> Vec<i64>;
214
215    fn glwe_trace_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
216    where
217        R: GLWEInfos,
218        A: GLWEInfos,
219        K: GGLWEInfos;
220
221    fn glwe_trace<R, A, K, H>(&self, res: &mut R, skip: usize, a: &A, keys: &H, scratch: &mut Scratch<BE>)
222    where
223        R: GLWEToMut + GLWEInfos,
224        A: GLWEToRef + GLWEInfos,
225        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
226        H: GLWEAutomorphismKeyHelper<K, BE>;
227
228    fn glwe_trace_inplace<R, K, H>(&self, res: &mut R, skip: usize, keys: &H, scratch: &mut Scratch<BE>)
229    where
230        R: GLWEToMut,
231        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
232        H: GLWEAutomorphismKeyHelper<K, BE>;
233}