1use poulpy_hal::{
2 api::{
3 ScratchTakeBasic, VecZnxAutomorphismInplace, VecZnxBigAutomorphismInplace, VecZnxBigSubSmallInplace,
4 VecZnxBigSubSmallNegateInplace,
5 },
6 layouts::{Backend, DataMut, Module, Scratch, VecZnxBig},
7};
8
9use crate::{
10 GLWEKeyswitch, ScratchTakeCore, keyswitch_internal,
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>
105where
106 Self: GLWEKeyswitch<BE>
107 + VecZnxAutomorphismInplace<BE>
108 + VecZnxBigAutomorphismInplace<BE>
109 + VecZnxBigSubSmallInplace<BE>
110 + VecZnxBigSubSmallNegateInplace<BE>,
111{
112 fn glwe_automorphism_tmp_bytes<R, A, K>(&self, res_infos: &R, a_infos: &A, key_infos: &K) -> usize
113 where
114 R: GLWEInfos,
115 A: GLWEInfos,
116 K: GGLWEInfos,
117 {
118 self.glwe_keyswitch_tmp_bytes(res_infos, a_infos, key_infos)
119 }
120
121 fn glwe_automorphism<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
122 where
123 R: GLWEToMut,
124 A: GLWEToRef,
125 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
126 Scratch<BE>: ScratchTakeCore<BE>,
127 {
128 self.glwe_keyswitch(res, a, key, scratch);
129
130 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
131
132 for i in 0..res.rank().as_usize() + 1 {
133 self.vec_znx_automorphism_inplace(key.p(), res.data_mut(), i, scratch);
134 }
135 }
136
137 fn glwe_automorphism_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
138 where
139 R: GLWEToMut,
140 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
141 Scratch<BE>: ScratchTakeCore<BE>,
142 {
143 self.glwe_keyswitch_inplace(res, key, scratch);
144
145 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
146
147 for i in 0..res.rank().as_usize() + 1 {
148 self.vec_znx_automorphism_inplace(key.p(), res.data_mut(), i, scratch);
149 }
150 }
151
152 fn glwe_automorphism_add<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
153 where
154 R: GLWEToMut,
155 A: GLWEToRef,
156 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
157 Scratch<BE>: ScratchTakeCore<BE>,
158 {
159 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
160 let a: &GLWE<&[u8]> = &a.to_ref();
161
162 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, key, scratch_1);
164
165 for i in 0..res.rank().as_usize() + 1 {
166 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
167 self.vec_znx_big_add_small_inplace(&mut res_big, i, a.data(), i);
168 self.vec_znx_big_normalize(
169 res.base2k().into(),
170 res.data_mut(),
171 i,
172 key.base2k().into(),
173 &res_big,
174 i,
175 scratch_1,
176 );
177 }
178 }
179
180 fn glwe_automorphism_add_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
181 where
182 R: GLWEToMut,
183 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
184 Scratch<BE>: ScratchTakeCore<BE>,
185 {
186 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
187
188 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, key, scratch_1);
190
191 for i in 0..res.rank().as_usize() + 1 {
192 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
193 self.vec_znx_big_add_small_inplace(&mut res_big, i, res.data(), i);
194 self.vec_znx_big_normalize(
195 res.base2k().into(),
196 res.data_mut(),
197 i,
198 key.base2k().into(),
199 &res_big,
200 i,
201 scratch_1,
202 );
203 }
204 }
205
206 fn glwe_automorphism_sub<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
207 where
208 R: GLWEToMut,
209 A: GLWEToRef,
210 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
211 Scratch<BE>: ScratchTakeCore<BE>,
212 {
213 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
214 let a: &GLWE<&[u8]> = &a.to_ref();
215
216 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, key, scratch_1);
218
219 for i in 0..res.rank().as_usize() + 1 {
220 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
221 self.vec_znx_big_sub_small_inplace(&mut res_big, i, a.data(), i);
222 self.vec_znx_big_normalize(
223 res.base2k().into(),
224 res.data_mut(),
225 i,
226 key.base2k().into(),
227 &res_big,
228 i,
229 scratch_1,
230 );
231 }
232 }
233
234 fn glwe_automorphism_sub_negate<R, A, K>(&self, res: &mut R, a: &A, key: &K, scratch: &mut Scratch<BE>)
235 where
236 R: GLWEToMut,
237 A: GLWEToRef,
238 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
239 Scratch<BE>: ScratchTakeCore<BE>,
240 {
241 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
242 let a: &GLWE<&[u8]> = &a.to_ref();
243
244 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, a, key, scratch_1);
246
247 for i in 0..res.rank().as_usize() + 1 {
248 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
249 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, a.data(), i);
250 self.vec_znx_big_normalize(
251 res.base2k().into(),
252 res.data_mut(),
253 i,
254 key.base2k().into(),
255 &res_big,
256 i,
257 scratch_1,
258 );
259 }
260 }
261
262 fn glwe_automorphism_sub_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
263 where
264 R: GLWEToMut,
265 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
266 Scratch<BE>: ScratchTakeCore<BE>,
267 {
268 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
269
270 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, key, scratch_1);
272
273 for i in 0..res.rank().as_usize() + 1 {
274 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
275 self.vec_znx_big_sub_small_inplace(&mut res_big, i, res.data(), i);
276 self.vec_znx_big_normalize(
277 res.base2k().into(),
278 res.data_mut(),
279 i,
280 key.base2k().into(),
281 &res_big,
282 i,
283 scratch_1,
284 );
285 }
286 }
287
288 fn glwe_automorphism_sub_negate_inplace<R, K>(&self, res: &mut R, key: &K, scratch: &mut Scratch<BE>)
289 where
290 R: GLWEToMut,
291 K: GetGaloisElement + GGLWEPreparedToRef<BE> + GGLWEInfos,
292 Scratch<BE>: ScratchTakeCore<BE>,
293 {
294 let res: &mut GLWE<&mut [u8]> = &mut res.to_mut();
295
296 let (res_dft, scratch_1) = scratch.take_vec_znx_dft(self, (res.rank() + 1).into(), key.size()); let mut res_big: VecZnxBig<_, BE> = keyswitch_internal(self, res_dft, res, key, scratch_1);
298
299 for i in 0..res.rank().as_usize() + 1 {
300 self.vec_znx_big_automorphism_inplace(key.p(), &mut res_big, i, scratch_1);
301 self.vec_znx_big_sub_small_negate_inplace(&mut res_big, i, res.data(), i);
302 self.vec_znx_big_normalize(
303 res.base2k().into(),
304 res.data_mut(),
305 i,
306 key.base2k().into(),
307 &res_big,
308 i,
309 scratch_1,
310 );
311 }
312 }
313}
314
315impl<BE: Backend> GLWEAutomorphism<BE> for Module<BE> where
316 Self: GLWEKeyswitch<BE>
317 + VecZnxAutomorphismInplace<BE>
318 + VecZnxBigAutomorphismInplace<BE>
319 + VecZnxBigSubSmallInplace<BE>
320 + VecZnxBigSubSmallNegateInplace<BE>
321{
322}