poulpy/poulpy-core/src/automorphism/glwe_ct.rs
poulpy-core/src/automorphism/glwe_ct.rsRaw
1use poulpy_hal::{
2 api::{
3 ScratchTakeBasic, VecZnxAutomorphismInplace, VecZnxBigAddSmallInplace, VecZnxBigAutomorphismInplace, VecZnxBigNormalize,
4 VecZnxBigSubSmallInplace, VecZnxBigSubSmallNegateInplace, VecZnxNormalize,
5 },
6 layouts::{Backend, DataMut, Module, Scratch, VecZnxBig},
7};
8
9use crate::{
10 GLWEKeySwitchInternal, GLWEKeyswitch, ScratchTakeCore,
11 layouts::{GGLWEInfos, GGLWEPreparedToRef, GLWE, GLWEInfos, GLWEToMut, GLWEToRef, GetGaloisElement, LWEInfos},
12};
13
14impl GLWE<Vec<u8>> {
15 pub fn automorphism_tmp_bytes<M, R, A, K, BE: Backend>(module: &M, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
16 where
17 R: GLWEInfos,
18 A: GLWEInfos,
19 K: GGLWEInfos,
20 M: GLWEAutomorphism<BE>,
21 {
22 module.glwe_automorphism_tmp_bytes(res_infos, a_infos, key_infos)
23 }
24}
25
26impl<DataSelf: DataMut> GLWE<DataSelf> {
27 pub fn automorphism<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
28 where
29 M: GLWEAutomorphism<BE>,
30 A: GLWEToRef,
31 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
32 Scratch<BE>: ScratchTakeCore<BE>,
33 {
34 module.glwe_automorphism(self, a, key, scratch);
35 }
36
37 pub fn automorphism_add<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
38 where
39 M: GLWEAutomorphism<BE>,
40 A: GLWEToRef,
41 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
42 Scratch<BE>: ScratchTakeCore<BE>,
43 {
44 module.glwe_automorphism_add(self, a, key, scratch);
45 }
46
47 pub fn automorphism_sub<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
48 where
49 M: GLWEAutomorphism<BE>,
50 A: GLWEToRef,
51 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
52 Scratch<BE>: ScratchTakeCore<BE>,
53 {
54 module.glwe_automorphism_sub(self, a, key, scratch);
55 }
56
57 pub fn automorphism_sub_negate<M, A, K, BE: Backend>(&mut self, module: &M, a: &A, key: &K, scratch: &mut Scratch<BE>)
58 where
59 M: GLWEAutomorphism<BE>,
60 A: GLWEToRef,
61 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
62 Scratch<BE>: ScratchTakeCore<BE>,
63 {
64 module.glwe_automorphism_sub_negate(self, a, key, scratch);
65 }
66
67 pub fn automorphism_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
68 where
69 M: GLWEAutomorphism<BE>,
70 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
71 Scratch<BE>: ScratchTakeCore<BE>,
72 {
73 module.glwe_automorphism_inplace(self, key, scratch);
74 }
75
76 pub fn automorphism_add_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
77 where
78 M: GLWEAutomorphism<BE>,
79 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
80 Scratch<BE>: ScratchTakeCore<BE>,
81 {
82 module.glwe_automorphism_add_inplace(self, key, scratch);
83 }
84
85 pub fn automorphism_sub_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
86 where
87 M: GLWEAutomorphism<BE>,
88 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
89 Scratch<BE>: ScratchTakeCore<BE>,
90 {
91 module.glwe_automorphism_sub_inplace(self, key, scratch);
92 }
93
94 pub fn automorphism_sub_negate_inplace<M, K, BE: Backend>(&mut self, module: &M, key: &K, scratch: &mut Scratch<BE>)
95 where
96 M: GLWEAutomorphism<BE>,
97 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
98 Scratch<BE>: ScratchTakeCore<BE>,
99 {
100 module.glwe_automorphism_sub_negate_inplace(self, key, scratch);
101 }
102}
103
104pub trait GLWEAutomorphism<BE: Backend> {
105 fn glwe_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
106 where
107 R: GLWEInfos,
108 A: GLWEInfos,
109 K: GGLWEInfos;
110
111 fn glwe_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
112 where
113 R: GLWEToMut,
114 A: GLWEToRef,
115 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
116
117 fn glwe_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
118 where
119 R: GLWEToMut,
120 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
121
122 fn glwe_automorphism_add<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
123 where
124 R: GLWEToMut,
125 A: GLWEToRef,
126 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
127
128 fn glwe_automorphism_add_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
129 where
130 R: GLWEToMut,
131 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
132
133 fn glwe_automorphism_sub<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
134 where
135 R: GLWEToMut,
136 A: GLWEToRef,
137 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
138
139 fn glwe_automorphism_sub_negate<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
140 where
141 R: GLWEToMut,
142 A: GLWEToRef,
143 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
144
145 fn glwe_automorphism_sub_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
146 where
147 R: GLWEToMut,
148 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
149
150 fn glwe_automorphism_sub_negate_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
151 where
152 R: GLWEToMut,
153 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos;
154}
155
156impl<BE: Backend> GLWEAutomorphism<BE> for Module<BE>
157where
158 Self: Sized
159 + GLWEKeyswitch<BE>
160 + GLWEKeySwitchInternal<BE>
161 + VecZnxNormalize<BE>
162 + VecZnxAutomorphismInplace<BE>
163 + VecZnxBigAutomorphismInplace<BE>
164 + VecZnxBigSubSmallInplace<BE>
165 + VecZnxBigSubSmallNegateInplace<BE>
166 + VecZnxBigAddSmallInplace<BE>
167 + VecZnxBigNormalize<BE>,
168 Scratch<BE>: ScratchTakeCore<BE>,
169{
170 fn glwe_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
171 where
172 R: GLWEInfos,
173 A: GLWEInfos,
174 K: GGLWEInfos,
175 {
176 self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
177 }
178
179 fn glwe_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
180 where
181 R: GLWEToMut,
182 A: GLWEToRef,
183 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
184 Scratch<BE>: ScratchTakeCore<BE>,
185 {
186 self.glwe_keyswitch(res, a, key, scratch);
187
188 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
189
190 for i in 0..res.rank().as_usize() + 1 {
191 self.vec_znx_automorphism_inplace(key.p(), res.data_mut(), i, scratch);
192 }
193 }
194
195 fn glwe_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
196 where
197 R: GLWEToMut,
198 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
199 Scratch<BE>: ScratchTakeCore<BE>,
200 {
201 self.glwe_keyswitch_inplace(res, key, scratch);
202
203 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
204
205 for i in 0..res.rank().as_usize() + 1 {
206 self.vec_znx_automorphism_inplace(key.p(), res.data_mut(), i, scratch);
207 }
208 }
209
210 fn glwe_automorphism_add<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
211 where
212 R: GLWEToMut,
213 A: GLWEToRef,
214 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
215 Scratch<BE>: ScratchTakeCore<BE>,
216 {
217 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
218 let a: &GLWE<&[u8]> = &a.to_ref();
219
220 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
221 let mut res_big: VecZnxBig<_, BE> = self.glwe_keyswitch_internal(res_dft, a, key, scratch_1);
222
223 for i in 0..res.rank().as_usize() + 1 {
224 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
225 self.vec_znx_big_add_small_inplace(&mut res_big, i, a.data(), i);
226 self.vec_znx_big_normalize(
227 res.base2k().into(),
228 res.data_mut(),
229 i,
230 key.base2k().into(),
231 &res_big,
232 i,
233 scratch_1,
234 );
235 }
236 }
237
238 fn glwe_automorphism_add_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
239 where
240 R: GLWEToMut,
241 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
242 Scratch<BE>: ScratchTakeCore<BE>,
243 {
244 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
245
246 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
247 let mut res_big: VecZnxBig<_, BE> = self.glwe_keyswitch_internal(res_dft, res, key, scratch_1);
248
249 for i in 0..res.rank().as_usize() + 1 {
250 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
251 self.vec_znx_big_add_small_inplace(&mut res_big, i, res.data(), i);
252 self.vec_znx_big_normalize(
253 res.base2k().into(),
254 res.data_mut(),
255 i,
256 key.base2k().into(),
257 &res_big,
258 i,
259 scratch_1,
260 );
261 }
262 }
263
264 fn glwe_automorphism_sub<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
265 where
266 R: GLWEToMut,
267 A: GLWEToRef,
268 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
269 Scratch<BE>: ScratchTakeCore<BE>,
270 {
271 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
272 let a: &GLWE<&[u8]> = &a.to_ref();
273
274 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
275 let mut res_big: VecZnxBig<_, BE> = self.glwe_keyswitch_internal(res_dft, a, key, scratch_1);
276
277 for i in 0..res.rank().as_usize() + 1 {
278 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
279 self.vec_znx_big_sub_small_inplace(&mut res_big, i, a.data(), i);
280 self.vec_znx_big_normalize(
281 res.base2k().into(),
282 res.data_mut(),
283 i,
284 key.base2k().into(),
285 &res_big,
286 i,
287 scratch_1,
288 );
289 }
290 }
291
292 fn glwe_automorphism_sub_negate<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
293 where
294 R: GLWEToMut,
295 A: GLWEToRef,
296 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
297 Scratch<BE>: ScratchTakeCore<BE>,
298 {
299 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
300 let a: &GLWE<&[u8]> = &a.to_ref();
301
302 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
303 let mut res_big: VecZnxBig<_, BE> = self.glwe_keyswitch_internal(res_dft, a, key, scratch_1);
304
305 for i in 0..res.rank().as_usize() + 1 {
306 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
307 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, a.data(), i);
308 self.vec_znx_big_normalize(
309 res.base2k().into(),
310 res.data_mut(),
311 i,
312 key.base2k().into(),
313 &res_big,
314 i,
315 scratch_1,
316 );
317 }
318 }
319
320 fn glwe_automorphism_sub_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
321 where
322 R: GLWEToMut,
323 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
324 Scratch<BE>: ScratchTakeCore<BE>,
325 {
326 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
327
328 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
329 let mut res_big: VecZnxBig<_, BE> = self.glwe_keyswitch_internal(res_dft, res, key, scratch_1);
330
331 for i in 0..res.rank().as_usize() + 1 {
332 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
333 self.vec_znx_big_sub_small_inplace(&mut res_big, i, res.data(), i);
334 self.vec_znx_big_normalize(
335 res.base2k().into(),
336 res.data_mut(),
337 i,
338 key.base2k().into(),
339 &res_big,
340 i,
341 scratch_1,
342 );
343 }
344 }
345
346 fn glwe_automorphism_sub_negate_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
347 where
348 R: GLWEToMut,
349 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
350 Scratch<BE>: ScratchTakeCore<BE>,
351 {
352 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
353
354 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); // TODO: optimise size
355 let mut res_big: VecZnxBig<_, BE> = self.glwe_keyswitch_internal(res_dft, res, key, scratch_1);
356
357 for i in 0..res.rank().as_usize() + 1 {
358 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
359 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, res.data(), i);
360 self.vec_znx_big_normalize(
361 res.base2k().into(),
362 res.data_mut(),
363 i,
364 key.base2k().into(),
365 &res_big,
366 i,
367 scratch_1,
368 );
369 }
370 }
371}