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