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>(
 35        &mut self,
 36        module: &M,
 37        start: usize,
 38        end: usize,
 39        a: &A,
 40        keys: &H,
 41        scratch: &mut Scratch<BE>,
 42    ) where
 43        A: GLWEToRef,
 44        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
 45        H: GLWEAutomorphismKeyHelper<K, BE>,
 46        Scratch<BE>: ScratchTakeCore<BE>,
 47        M: GLWETrace<BE>,
 48    {
 49        module.glwe_trace(self, start, end, a, keys, scratch);
 50    }
 51
 52    pub fn trace_inplace<H, K, M, BE: Backend>(
 53        &mut self,
 54        module: &M,
 55        start: usize,
 56        end: usize,
 57        keys: &H,
 58        scratch: &mut Scratch<BE>,
 59    ) where
 60        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
 61        H: GLWEAutomorphismKeyHelper<K, BE>,
 62        Scratch<BE>: ScratchTakeCore<BE>,
 63        M: GLWETrace<BE>,
 64    {
 65        module.glwe_trace_inplace(self, start, end, keys, scratch);
 66    }
 67}
 68
 69#[inline(always)]
 70pub fn trace_galois_elements(log_n: usize, cyclotomic_order: i64) -> Vec<i64> {
 71    (0..log_n)
 72        .map(|i| {
 73            if i == 0 {
 74                -1
 75            } else {
 76                galois_element(1 << (i - 1), cyclotomic_order)
 77            }
 78        })
 79        .collect()
 80}
 81
 82impl<BE: Backend> GLWETrace<BE> for Module<BE>
 83where
 84    Self: ModuleLogN
 85        + GaloisElement
 86        + GLWEAutomorphism<BE>
 87        + GLWEShift<BE>
 88        + GLWECopy
 89        + CyclotomicOrder
 90        + VecZnxNormalizeTmpBytes
 91        + VecZnxNormalize<BE>,
 92    Scratch<BE>: ScratchTakeCore<BE>,
 93{
 94    fn glwe_trace_galois_elements(&self) -> Vec<i64> {
 95        trace_galois_elements(self.log_n(), self.cyclotomic_order())
 96    }
 97
 98    fn glwe_trace_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
 99    where
100        R: GLWEInfos,
101        A: GLWEInfos,
102        K: GGLWEInfos,
103    {
104        let trace: usize = self.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos);
105        if a_infos.base2k() != key_infos.base2k() {
106            let glwe_conv: usize = VecZnx::bytes_of(
107                self.n(),
108                (key_infos.rank_out() + 1).into(),
109                res_infos.k().min(a_infos.k()).div_ceil(key_infos.base2k()) as usize,
110            ) + self.vec_znx_normalize_tmp_bytes();
111            return glwe_conv + trace;
112        }
113
114        trace
115    }
116
117    fn glwe_trace<R, A, K, H>(&self, res: &mut R, start: usize, end: usize, a: &A, keys: &H, scratch: &mut Scratch<BE>)
118    where
119        R: GLWEToMut,
120        A: GLWEToRef,
121        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
122        H: GLWEAutomorphismKeyHelper<K, BE>,
123    {
124        self.glwe_copy(res, a);
125        self.glwe_trace_inplace(res, start, end, keys, scratch);
126    }
127
128    fn glwe_trace_inplace<R, K, H>(&self, res: &mut R, start: usize, end: 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
138        assert_eq!(res.n(), self.n() as u32);
139        assert_eq!(ksk_infos.n(), self.n() as u32);
140        assert!(start < end);
141        assert!(end <= self.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 start..end {
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            // println!("res: {}", res);
194
195            for i in start..end {
196                self.glwe_rsh(1, res, scratch);
197
198                let p: i64 = if i == 0 {
199                    -1
200                } else {
201                    self.galois_element(1 << (i - 1))
202                };
203
204                if let Some(key) = keys.get_automorphism_key(p) {
205                    self.glwe_automorphism_add_inplace(res, key, scratch);
206                } else {
207                    panic!("keys[{p}] is empty")
208                }
209            }
210        }
211    }
212}
213
214pub trait GLWETrace<BE: Backend> {
215    fn glwe_trace_galois_elements(&self) -> Vec<i64>;
216
217    fn glwe_trace_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
218    where
219        R: GLWEInfos,
220        A: GLWEInfos,
221        K: GGLWEInfos;
222
223    fn glwe_trace<R, A, K, H>(&self, res: &mut R, start: usize, end: usize, a: &A, keys: &H, scratch: &mut Scratch<BE>)
224    where
225        R: GLWEToMut,
226        A: GLWEToRef,
227        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
228        H: GLWEAutomorphismKeyHelper<K, BE>;
229
230    fn glwe_trace_inplace<R, K, H>(&self, res: &mut R, start: usize, end: usize, keys: &H, scratch: &mut Scratch<BE>)
231    where
232        R: GLWEToMut,
233        K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
234        H: GLWEAutomorphismKeyHelper<K, BE>;
235}