1use poulpy_hal::{
  2    api::{ModuleLogN, VecZnxNormalizeTmpBytes},
  3    layouts::{Backend, CyclotomicOrder, DataMut, GaloisElement, Module, Scratch, VecZnx, galois_element},
  4};
  5
  6use crate::{
  7    GLWEAutomorphism, GLWECopy, GLWENormalize, 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        + GLWENormalize<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 atk_layout: &GGLWELayout = &keys.automorphism_key_infos();
118
119        let (mut tmp, scratch_1) = scratch.take_glwe(&GLWELayout {
120            n: res.n(),
121            base2k: atk_layout.base2k(),
122            k: a.k().max(res.k()),
123            rank: res.rank(),
124        });
125
126        if a.base2k() == atk_layout.base2k() {
127            self.glwe_copy(&mut tmp, a);
128        } else {
129            self.glwe_normalize(&mut tmp, a, scratch_1);
130        }
131
132        self.glwe_trace_inplace(&mut tmp, skip, keys, scratch_1);
133
134        if res.base2k() == atk_layout.base2k() {
135            self.glwe_copy(res, &tmp);
136        } else {
137            self.glwe_normalize(res, &tmp, scratch_1);
138        }
139    }
140
141    fn glwe_trace_inplace<R, K, H>(&self, res: &mut R, skip: usize, keys: &H, scratch: &mut Scratch<BE>)
142    where
143        R: GLWEToMut,
144        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
145        H: GLWEAutomorphismKeyHelper<K, BE>,
146    {
147        let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
148
149        let ksk_infos: &GGLWELayout = &keys.automorphism_key_infos();
150        let log_n: usize = self.log_n();
151
152        assert_eq!(res.n(), self.n() as u32);
153        assert_eq!(ksk_infos.n(), self.n() as u32);
154        assert!(skip <= log_n);
155        assert_eq!(ksk_infos.rank_in(), res.rank());
156        assert_eq!(ksk_infos.rank_out(), res.rank());
157
158        if res.base2k() != ksk_infos.base2k() {
159            let (mut res_conv, scratch_1) = scratch.take_glwe(&GLWELayout {
160                n: self.n().into(),
161                base2k: ksk_infos.base2k(),
162                k: res.k(),
163                rank: res.rank(),
164            });
165            self.glwe_normalize(&mut res_conv, res, scratch_1);
166            self.glwe_trace_inplace(&mut res_conv, skip, keys, scratch_1);
167            self.glwe_normalize(res, &res_conv, scratch_1);
168        } else {
169            for i in skip..log_n {
170                self.glwe_rsh(1, res, scratch);
171
172                let p: i64 = if i == 0 { -1 } else { self.galois_element(1 << (i - 1)) };
173
174                if let Some(key) = keys.get_automorphism_key(p) {
175                    self.glwe_automorphism_add_inplace(res, key, scratch);
176                } else {
177                    panic!("keys[{p}] is empty")
178                }
179            }
180        }
181    }
182}
183
184pub trait GLWETrace<BE: Backend> {
185    fn glwe_trace_galois_elements(&self) -> Vec<i64>;
186
187    fn glwe_trace_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
188    where
189        R: GLWEInfos,
190        A: GLWEInfos,
191        K: GGLWEInfos;
192
193    fn glwe_trace<R, A, K, H>(&self, res: &mut R, skip: usize, a: &A, keys: &H, scratch: &mut Scratch<BE>)
194    where
195        R: GLWEToMut + GLWEInfos,
196        A: GLWEToRef + GLWEInfos,
197        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
198        H: GLWEAutomorphismKeyHelper<K, BE>;
199
200    fn glwe_trace_inplace<R, K, H>(&self, res: &mut R, skip: usize, keys: &H, scratch: &mut Scratch<BE>)
201    where
202        R: GLWEToMut,
203        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
204        H: GLWEAutomorphismKeyHelper<K, BE>;
205}