poulpy/poulpy-core/src/glwe_trace.rs
poulpy-core/src/glwe_trace.rsRaw
1use std::collections::HashMap;
2
3use poulpy_hal::{
4 api::{ModuleLogN, VecZnxNormalize, VecZnxNormalizeTmpBytes},
5 layouts::{Backend, CyclotomicOrder, DataMut, GaloisElement, Module, Scratch, VecZnx, galois_element},
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_trace_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
68#[inline(always)]
69pub fn trace_galois_elements(log_n: usize, cyclotomic_order: i64) -> Vec<i64> {
70 (0..log_n)
71 .map(|i| {
72 if i == 0 {
73 -1
74 } else {
75 galois_element(1 << (i - 1), cyclotomic_order)
76 }
77 })
78 .collect()
79}
80
81impl<BE: Backend> GLWETrace<BE> for Module<BE>
82where
83 Self: ModuleLogN
84 + GaloisElement
85 + GLWEAutomorphism<BE>
86 + GLWEShift<BE>
87 + GLWECopy
88 + CyclotomicOrder
89 + VecZnxNormalizeTmpBytes
90 + VecZnxNormalize<BE>,
91 Scratch<BE>: ScratchTakeCore<BE>,
92{
93 fn glwe_trace_galois_elements(&self) -> Vec<i64> {
94 trace_galois_elements(self.log_n(), self.cyclotomic_order())
95 }
96
97 fn glwe_trace_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
98 where
99 R: GLWEInfos,
100 A: GLWEInfos,
101 K: GGLWEInfos,
102 {
103 let trace: usize = self.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos);
104 if a_infos.base2k() != key_infos.base2k() {
105 let glwe_conv: usize = VecZnx::bytes_of(
106 self.n(),
107 (key_infos.rank_out() + 1).into(),
108 res_infos.k().min(a_infos.k()).div_ceil(key_infos.base2k()) as usize,
109 ) + self.vec_znx_normalize_tmp_bytes();
110 return glwe_conv + trace;
111 }
112
113 trace
114 }
115
116 fn glwe_trace<R, A, K>(&self, res: &mut R, start: usize, end: usize, a: &A, keys: &HashMap<i64, K>, scratch: &mut Scratch<BE>)
117 where
118 R: GLWEToMut,
119 A: GLWEToRef,
120 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
121 {
122 self.glwe_copy(res, a);
123 self.glwe_trace_inplace(res, start, end, keys, scratch);
124 }
125
126 fn glwe_trace_inplace<R, K>(&self, res: &mut R, start: usize, end: usize, keys: &HashMap<i64, K>, scratch: &mut Scratch<BE>)
127 where
128 R: GLWEToMut,
129 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos,
130 {
131 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
132
133 let basek_ksk: Base2K = keys.get(keys.keys().next().unwrap()).unwrap().base2k();
134
135 #[cfg(debug_assertions)]
136 {
137 assert_eq!(res.n(), self.n() as u32);
138 assert!(start < end);
139 assert!(end <= self.log_n());
140 for key in keys.values() {
141 assert_eq!(key.n(), self.n() as u32);
142 assert_eq!(key.base2k(), basek_ksk);
143 assert_eq!(key.rank_in(), res.rank());
144 assert_eq!(key.rank_out(), res.rank());
145 }
146 }
147
148 if res.base2k() != basek_ksk {
149 let (mut self_conv, scratch_1) = scratch.take_glwe(&GLWELayout {
150 n: self.n().into(),
151 base2k: basek_ksk,
152 k: res.k(),
153 rank: res.rank(),
154 });
155
156 for j in 0..(res.rank() + 1).into() {
157 self.vec_znx_normalize(
158 basek_ksk.into(),
159 &mut self_conv.data,
160 j,
161 basek_ksk.into(),
162 res.data(),
163 j,
164 scratch_1,
165 );
166 }
167
168 for i in start..end {
169 self.glwe_rsh(1, &mut self_conv, scratch_1);
170
171 let p: i64 = if i == 0 {
172 -1
173 } else {
174 self.galois_element(1 << (i - 1))
175 };
176
177 if let Some(key) = keys.get(&p) {
178 self.glwe_automorphism_add_inplace(&mut self_conv, key, scratch_1);
179 } else {
180 panic!("keys[{p}] is empty")
181 }
182 }
183
184 for j in 0..(res.rank() + 1).into() {
185 self.vec_znx_normalize(
186 res.base2k().into(),
187 res.data_mut(),
188 j,
189 basek_ksk.into(),
190 &self_conv.data,
191 j,
192 scratch_1,
193 );
194 }
195 } else {
196 // println!("res: {}", res);
197
198 for i in start..end {
199 self.glwe_rsh(1, res, scratch);
200
201 let p: i64 = if i == 0 {
202 -1
203 } else {
204 self.galois_element(1 << (i - 1))
205 };
206
207 if let Some(key) = keys.get(&p) {
208 self.glwe_automorphism_add_inplace(res, key, scratch);
209 } else {
210 panic!("keys[{p}] is empty")
211 }
212 }
213 }
214 }
215}
216
217pub trait GLWETrace<BE: Backend> {
218 fn glwe_trace_galois_elements(&self) -> Vec<i64>;
219
220 fn glwe_trace_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
221 where
222 R: GLWEInfos,
223 A: GLWEInfos,
224 K: GGLWEInfos;
225
226 fn glwe_trace<R, A, K>(
227 &self,
228 res: &mut R,
229 start: usize,
230 end: usize,
231 a: &A,
232 keys: &HashMap<i64, K>,
233 scratch: &mut Scratch<BE>,
234 ) where
235 R: GLWEToMut,
236 A: GLWEToRef,
237 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos;
238
239 fn glwe_trace_inplace<R, K>(&self, res: &mut R, start: usize, end: usize, keys: &HashMap<i64, K>, scratch: &mut Scratch<BE>)
240 where
241 R: GLWEToMut,
242 K: GGLWEPreparedToRef<BE> + GetGaloisElement + GGLWEInfos;
243}