poulpy/core/src/ggsw_ciphertext.rs
core/src/ggsw_ciphertext.rsRaw
1use base2k::{
2 Backend, FFT64, MatZnxDft, MatZnxDftAlloc, MatZnxDftOps, MatZnxDftToMut, MatZnxDftToRef, Module, ScalarZnx, ScalarZnxDft,
3 ScalarZnxDftToRef, ScalarZnxToRef, Scratch, VecZnxAlloc, VecZnxBigOps, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut,
4 VecZnxDftToRef, VecZnxOps, ZnxInfos, ZnxZero,
5};
6use sampling::source::Source;
7
8use crate::{
9 automorphism::AutomorphismKey,
10 elem::{GetRow, Infos, SetRow},
11 glwe_ciphertext::GLWECiphertext,
12 glwe_ciphertext_fourier::GLWECiphertextFourier,
13 glwe_plaintext::GLWEPlaintext,
14 keys::SecretKeyFourier,
15 keyswitch_key::GLWESwitchingKey,
16 tensor_key::TensorKey,
17 utils::derive_size,
18};
19
20pub struct GGSWCiphertext<C, B: Backend> {
21 pub data: MatZnxDft<C, B>,
22 pub basek: usize,
23 pub k: usize,
24}
25
26impl<B: Backend> GGSWCiphertext<Vec<u8>, B> {
27 pub fn new(module: &Module<B>, basek: usize, k: usize, rows: usize, rank: usize) -> Self {
28 Self {
29 data: module.new_mat_znx_dft(rows, rank + 1, rank + 1, derive_size(basek, k)),
30 basek: basek,
31 k: k,
32 }
33 }
34}
35
36impl<T, B: Backend> Infos for GGSWCiphertext<T, B> {
37 type Inner = MatZnxDft<T, B>;
38
39 fn inner(&self) -> &Self::Inner {
40 &self.data
41 }
42
43 fn basek(&self) -> usize {
44 self.basek
45 }
46
47 fn k(&self) -> usize {
48 self.k
49 }
50}
51
52impl<T, B: Backend> GGSWCiphertext<T, B> {
53 pub fn rank(&self) -> usize {
54 self.data.cols_out() - 1
55 }
56}
57
58impl<C, B: Backend> MatZnxDftToMut<B> for GGSWCiphertext<C, B>
59where
60 MatZnxDft<C, B>: MatZnxDftToMut<B>,
61{
62 fn to_mut(&mut self) -> MatZnxDft<&mut [u8], B> {
63 self.data.to_mut()
64 }
65}
66
67impl<C, B: Backend> MatZnxDftToRef<B> for GGSWCiphertext<C, B>
68where
69 MatZnxDft<C, B>: MatZnxDftToRef<B>,
70{
71 fn to_ref(&self) -> MatZnxDft<&[u8], B> {
72 self.data.to_ref()
73 }
74}
75
76impl GGSWCiphertext<Vec<u8>, FFT64> {
77 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, rank: usize, size: usize) -> usize {
78 GLWECiphertext::encrypt_sk_scratch_space(module, size)
79 + module.bytes_of_vec_znx(rank + 1, size)
80 + module.bytes_of_vec_znx(1, size)
81 + module.bytes_of_vec_znx_dft(rank + 1, size)
82 }
83
84 pub fn automorphism_scratch_space(
85 module: &Module<FFT64>,
86 out_size: usize,
87 in_size: usize,
88 auto_key_size: usize,
89 rank: usize,
90 ) -> usize {
91 let tmp_dft: usize = module.bytes_of_vec_znx_dft(rank + 1, auto_key_size);
92 let tmp_idft: usize = module.bytes_of_vec_znx(rank + 1, out_size);
93 let vmp: usize = GLWECiphertext::keyswitch_from_fourier_scratch_space(module, out_size, rank, in_size, rank, auto_key_size);
94 tmp_dft + tmp_idft + vmp
95 }
96
97 pub fn external_product_scratch_space(
98 module: &Module<FFT64>,
99 out_size: usize,
100 in_size: usize,
101 ggsw_size: usize,
102 rank: usize,
103 ) -> usize {
104 let tmp_in: usize = module.bytes_of_vec_znx_dft(rank + 1, in_size);
105 let tmp_out: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
106 let ggsw: usize = GLWECiphertextFourier::external_product_scratch_space(module, out_size, in_size, ggsw_size, rank);
107 tmp_in + tmp_out + ggsw
108 }
109
110 pub fn external_product_inplace_scratch_space(
111 module: &Module<FFT64>,
112 out_size: usize,
113 ggsw_size: usize,
114 rank: usize,
115 ) -> usize {
116 let tmp: usize = module.bytes_of_vec_znx_dft(rank + 1, out_size);
117 let ggsw: usize = GLWECiphertextFourier::external_product_inplace_scratch_space(module, out_size, ggsw_size, rank);
118 tmp + ggsw
119 }
120}
121
122impl<DataSelf> GGSWCiphertext<DataSelf, FFT64>
123where
124 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64> + MatZnxDftToRef<FFT64>,
125{
126 pub fn encrypt_sk<DataPt, DataSk>(
127 &mut self,
128 module: &Module<FFT64>,
129 pt: &ScalarZnx<DataPt>,
130 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
131 source_xa: &mut Source,
132 source_xe: &mut Source,
133 sigma: f64,
134 scratch: &mut Scratch,
135 ) where
136 ScalarZnx<DataPt>: ScalarZnxToRef,
137 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
138 {
139 #[cfg(debug_assertions)]
140 {
141 assert_eq!(self.rank(), sk_dft.rank());
142 assert_eq!(self.n(), module.n());
143 assert_eq!(pt.n(), module.n());
144 assert_eq!(sk_dft.n(), module.n());
145 }
146
147 let size: usize = self.size();
148 let basek: usize = self.basek();
149 let k: usize = self.k();
150 let cols: usize = self.rank() + 1;
151
152 let (tmp_znx_pt, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
153 let (tmp_znx_ct, scrach_2) = scratch_1.tmp_vec_znx(module, cols, size);
154
155 let mut vec_znx_pt: GLWEPlaintext<&mut [u8]> = GLWEPlaintext {
156 data: tmp_znx_pt,
157 basek: basek,
158 k: k,
159 };
160
161 let mut vec_znx_ct: GLWECiphertext<&mut [u8]> = GLWECiphertext {
162 data: tmp_znx_ct,
163 basek: basek,
164 k,
165 };
166
167 (0..self.rows()).for_each(|row_j| {
168 vec_znx_pt.data.zero();
169
170 // Adds the scalar_znx_pt to the i-th limb of the vec_znx_pt
171 module.vec_znx_add_scalar_inplace(&mut vec_znx_pt, 0, row_j, pt, 0);
172 module.vec_znx_normalize_inplace(basek, &mut vec_znx_pt, 0, scrach_2);
173
174 (0..cols).for_each(|col_i| {
175 // rlwe encrypt of vec_znx_pt into vec_znx_ct
176
177 vec_znx_ct.encrypt_sk_private(
178 module,
179 Some((&vec_znx_pt, col_i)),
180 sk_dft,
181 source_xa,
182 source_xe,
183 sigma,
184 scrach_2,
185 );
186
187 // Switch vec_znx_ct into DFT domain
188 {
189 let (mut vec_znx_dft_ct, _) = scrach_2.tmp_vec_znx_dft(module, cols, size);
190
191 (0..cols).for_each(|i| {
192 module.vec_znx_dft(&mut vec_znx_dft_ct, i, &vec_znx_ct, i);
193 });
194
195 module.vmp_prepare_row(self, row_j, col_i, &vec_znx_dft_ct);
196 }
197 });
198 });
199 }
200
201
202 pub fn keyswitch<DataLhs, DataRhs0, DataRhs1>(
203 &mut self,
204 module: &Module<FFT64>,
205 lhs: &GGSWCiphertext<DataLhs, FFT64>,
206 ksk: &GLWESwitchingKey<DataRhs0, FFT64>,
207 tsk: &TensorKey<DataRhs1, FFT64>,
208 scratch: &mut Scratch,
209 ) where
210 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
211 MatZnxDft<DataRhs0, FFT64>: MatZnxDftToRef<FFT64>,
212 MatZnxDft<DataRhs1, FFT64>: MatZnxDftToRef<FFT64>,
213 {
214 #[cfg(debug_assertions)]
215 {
216 assert_eq!(
217 self.rank(),
218 lhs.rank(),
219 "ggsw_out rank: {} != ggsw_in rank: {}",
220 self.rank(),
221 lhs.rank()
222 );
223 assert_eq!(
224 lhs.rank(),
225 ksk.rank(),
226 "ggsw_in rank: {} != ksk rank: {}",
227 lhs.rank(),
228 ksk.rank()
229 );
230 assert_eq!(
231 lhs.rank(),
232 tsk.rank(),
233 "ggsw_in rank: {} != tsk rank: {}",
234 lhs.rank(),
235 tsk.rank()
236 );
237 }
238
239 let cols: usize = self.rank() + 1;
240
241 // Example for rank 3:
242 //
243 // Note: M is a vector (m, Bm, B^2m, B^3m, ...), so each column is
244 // actually composed of that many rows.
245 //
246 // # Input
247 //
248 // col 0: (-(a0s0 + a1s1 + a2s2) + M, a0 , a1 , a2 )
249 // col 1: (-(b0s0 + b1s1 + b2s2) , b0 + M, b1 , b2 )
250 // col 2: (-(c0s0 + c1s1 + c2s2) , c0 , c1 + M, c2 )
251 // col 3: (-(d0s0 + d1s1 + d2s2) , d0 , d1 , d2 + M)
252 //
253 // # Output
254 //
255 // col 0: (-(a0s0' + a1s1' + a2s2') + M, a0 , a1 , a2 )
256 // col 1: (-(b0s0' + b1s1' + b2s2') , b0 + M, b1 , b2 )
257 // col 2: (-(c0s0' + c1s1' + c2s2') , c0 , c1 + M, c2 )
258 // col 3: (-(d0s0' + d1s1' + d2s2') , d0 , d1 , d2 + M)
259 (0..self.rows()).for_each(|row_j| {
260
261 let (tmp_dft_out_data, scratch1) = scratch.tmp_vec_znx_dft(module, self.rank() + 1, self.size());
262
263 let mut tmp_dft_out: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
264 data: tmp_dft_out_data,
265 basek: lhs.basek(),
266 k: lhs.k(),
267 };
268
269 {
270 let (tmp_dft_in_data, scratch2) = scratch1.tmp_vec_znx_dft(module, lhs.rank() + 1, lhs.size());
271
272 let mut tmp_dft_in: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
273 data: tmp_dft_in_data,
274 basek: lhs.basek(),
275 k: lhs.k(),
276 };
277
278 // 1) Applies key-switching to GGSW[i][0]: (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2)
279 lhs.get_row(module, row_j, 0, &mut tmp_dft_in);
280 // (-(a0s0 + a1s1 + a2s2) + M[i], a0, a1, a2) -> (-(a0s0' + a1s1' + a2s2') + M[i], a0, a1, a2)
281 tmp_dft_out.keyswitch(module, &tmp_dft_in, ksk, scratch2);
282 self.set_row(module, row_j, 0, &tmp_dft_out);
283 }
284
285 // 2) Isolates IDFT(-(a0s0' + a1s1' + a2s2') + M[i])
286 let (mut tmp_c0_data, scratch2) = scratch1.tmp_vec_znx_big(module, 1, self.size());
287 module.vec_znx_idft_tmp_a(&mut tmp_c0_data, 0, &mut tmp_dft_out, 0);
288
289 // 3) Expands the i-th row of the other columns using the tensor key
290 // col 1: (-(b0s0' + b1s1' + b2s2') , b0 + M[i], b1 , b2 ) = KS_{s0's0', s0's1', s0's2'}(a0) + (0, -(a0s0' + a1s1' + a2s2') + M[i], 0, 0)
291 // col 2: (-(c0s0' + c1s1' + c2s2') , c0 , c1 + M[i], c2 ) = KS_{s1's0', s1's1', s1's2'}(a1) + (0, 0, -(a0s0' + a1s1' + a2s2') + M[i], 0)
292 // col 3: (-(d0s0' + d1s1' + d2s2') , d0 , d1 , d2 + M[i]) = KS_{s2's0', s2's1', s2's2'}(a2) + (0, 0, 0, -(a0s0' + a1s1' + a2s2') + M[i])
293 (1..cols).for_each(|col_i| {
294
295 let (tmp_dft_i_data, scratch3) = scratch2.tmp_vec_znx_dft(module, cols, tsk.size());
296 let mut tmp_dft_i: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
297 data: tmp_dft_i_data,
298 basek: lhs.basek(),
299 k: lhs.k(),
300 };
301
302 // 5) Performs a key-switch for each combination of s[i]*s[j], i.e. for a0, a1, a2
303 //
304 // # Example for col=1
305 //
306 // a0 * (-(f0s0 + f1s1 + f1s2) + s0^2, f0, f1, f2) = (-(a0f0s0 + a0f1s1 + a0f1s2) + a0s0^2, a0f0, a0f1, a0f2)
307 // +
308 // a1 * (-(g0s0 + g1s1 + g1s2) + s0s1, g0, g1, g2) = (-(a1g0s0 + a1g1s1 + a1g1s2) + a1s0s1, a1g0, a1g1, a1g2)
309 // +
310 // a2 * (-(h0s0 + h1s1 + h1s2) + s0s2, h0, h1, h2) = (-(a2h0s0 + a2h1s1 + a2h1s2) + a2s0s2, a2h0, a2h1, a2h2)
311 // =
312 // (-(x0s0' + x1s1' + x2s2') + s0'(a0s0' + a1s1' + a2s2'), x0, x1, x2)
313 (1..cols).for_each(|col_j| {
314
315 // Extracts a[i] and multipies with Enc(s'[i]s'[j])
316 let (mut tmp_dft_col_data, scratch4) = scratch3.tmp_vec_znx_dft(module, 1, self.size());
317 tmp_dft_col_data.extract_column(0, &tmp_dft_out.data, col_j);
318
319 if col_j == 1 {
320 module.vmp_apply(
321 &mut tmp_dft_i,
322 &tmp_dft_col_data,
323 tsk.at(col_i - 1, col_j - 1), // Selects Enc(s'[i]s'[j])
324 scratch4,
325 );
326 } else {
327 module.vmp_apply_add(
328 &mut tmp_dft_i,
329 &tmp_dft_col_data,
330 tsk.at(col_i - 1, col_j - 1), // Selects Enc(s'[i]s'[j])
331 scratch4,
332 );
333 }
334 });
335
336 // Adds -(sum a[i] * s[i]) + m) on the i-th column of tmp_idft_i
337 //
338 // (-(x0s0' + x1s1' + x2s2') + a0s0's0' + a1s0's1' + a2s0's2', x0, x1, x2)
339 // +
340 // (0, -(a0s0' + a1s1' + a2s2') + M[i], 0, 0)
341 // =
342 // (-(x0s0' + x1s1' + x2s2') + s0'(a0s0' + a1s1' + a2s2'), x0 -(a0s0' + a1s1' + a2s2') + M[i], x1, x2)
343 // =
344 // (-(x0s0' + x1s1' + x2s2'), x0 + M[i], x1, x2)
345 {
346 let (mut tmp_idft, scratch3) = scratch3.tmp_vec_znx_big(module, 1, tsk.size());
347 let (mut tmp_znx_small, scratch5) = scratch3.tmp_vec_znx(module, 1, self.size());
348 (0..cols).for_each(|i| {
349 module.vec_znx_idft_tmp_a(&mut tmp_idft, 0, &mut tmp_dft_i, i);
350 module.vec_znx_big_add_inplace(&mut tmp_idft, col_i, &tmp_c0_data, 0);
351 module.vec_znx_big_normalize(self.basek(), &mut tmp_znx_small, 0, &tmp_idft, 0, scratch5);
352 module.vec_znx_dft(&mut tmp_dft_i, i, &tmp_znx_small, 0);
353 });
354 }
355
356 // Stores (-(x0s0' + x1s1' + x2s2'), x0 + M[i], x1, x2)
357 self.set_row(module, row_j, col_i, &tmp_dft_i);
358 })
359 })
360 }
361
362 pub fn automorphism<DataLhs, DataRhs>(
363 &mut self,
364 module: &Module<FFT64>,
365 lhs: &GGSWCiphertext<DataLhs, FFT64>,
366 rhs: &AutomorphismKey<DataRhs, FFT64>,
367 scratch: &mut Scratch,
368 ) where
369 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
370 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
371 {
372 #[cfg(debug_assertions)]
373 {
374 assert_eq!(
375 self.rank(),
376 lhs.rank(),
377 "ggsw_out rank: {} != ggsw_in rank: {}",
378 self.rank(),
379 lhs.rank()
380 );
381 assert_eq!(
382 self.rank(),
383 rhs.rank(),
384 "ggsw_in rank: {} != auto_key rank: {}",
385 self.rank(),
386 rhs.rank()
387 );
388 }
389;
390 let cols: usize = self.rank() + 1;
391
392 let (tmp_dft_data, scratch1) = scratch.tmp_vec_znx_dft(module, cols, rhs.size()); //TODO optimize
393
394 let mut tmp_dft: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
395 data: tmp_dft_data,
396 basek: lhs.basek(),
397 k: lhs.k(),
398 };
399
400 let (tmp_idft_data, scratch2) = scratch1.tmp_vec_znx(module, cols, self.size());
401
402 let mut tmp_idft: GLWECiphertext<&mut [u8]> = GLWECiphertext::<&mut [u8]> {
403 data: tmp_idft_data,
404 basek: self.basek(),
405 k: self.k(),
406 };
407
408 (0..cols).for_each(|col_i| {
409 (0..self.rows()).for_each(|row_j| {
410 lhs.get_row(module, row_j, col_i, &mut tmp_dft);
411 tmp_idft.keyswitch_from_fourier(module, &tmp_dft, &rhs.key, scratch2);
412 (0..cols).for_each(|i| {
413 module.vec_znx_automorphism_inplace(rhs.p(), &mut tmp_idft, i);
414 });
415 self.set_row(module, row_j, col_i, &tmp_dft);
416 });
417 });
418
419 tmp_dft.data.zero();
420
421 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
422 (0..self.rank() + 1).for_each(|col_j| {
423 self.set_row(module, row_i, col_j, &tmp_dft);
424 });
425 });
426 }
427
428 pub fn external_product<DataLhs, DataRhs>(
429 &mut self,
430 module: &Module<FFT64>,
431 lhs: &GGSWCiphertext<DataLhs, FFT64>,
432 rhs: &GGSWCiphertext<DataRhs, FFT64>,
433 scratch: &mut Scratch,
434 ) where
435 MatZnxDft<DataLhs, FFT64>: MatZnxDftToRef<FFT64>,
436 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
437 {
438 #[cfg(debug_assertions)]
439 {
440 assert_eq!(
441 self.rank(),
442 lhs.rank(),
443 "ggsw_out rank: {} != ggsw_in rank: {}",
444 self.rank(),
445 lhs.rank()
446 );
447 assert_eq!(
448 self.rank(),
449 rhs.rank(),
450 "ggsw_in rank: {} != ggsw_apply rank: {}",
451 self.rank(),
452 rhs.rank()
453 );
454 }
455
456 let (tmp_in_data, scratch1) = scratch.tmp_vec_znx_dft(module, lhs.rank() + 1, lhs.size());
457
458 let mut tmp_in: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
459 data: tmp_in_data,
460 basek: lhs.basek(),
461 k: lhs.k(),
462 };
463
464 let (tmp_out_data, scratch2) = scratch1.tmp_vec_znx_dft(module, self.rank() + 1, self.size());
465
466 let mut tmp_out: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
467 data: tmp_out_data,
468 basek: self.basek(),
469 k: self.k(),
470 };
471
472 (0..self.rank() + 1).for_each(|col_i| {
473 (0..self.rows()).for_each(|row_j| {
474 lhs.get_row(module, row_j, col_i, &mut tmp_in);
475 tmp_out.external_product(module, &tmp_in, rhs, scratch2);
476 self.set_row(module, row_j, col_i, &tmp_out);
477 });
478 });
479
480 tmp_out.data.zero();
481
482 (self.rows().min(lhs.rows())..self.rows()).for_each(|row_i| {
483 (0..self.rank() + 1).for_each(|col_j| {
484 self.set_row(module, row_i, col_j, &tmp_out);
485 });
486 });
487 }
488
489 pub fn external_product_inplace<DataRhs>(
490 &mut self,
491 module: &Module<FFT64>,
492 rhs: &GGSWCiphertext<DataRhs, FFT64>,
493 scratch: &mut Scratch,
494 ) where
495 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
496 {
497 #[cfg(debug_assertions)]
498 {
499 assert_eq!(
500 self.rank(),
501 rhs.rank(),
502 "ggsw_out rank: {} != ggsw_apply: {}",
503 self.rank(),
504 rhs.rank()
505 );
506 }
507
508 let (tmp_data, scratch1) = scratch.tmp_vec_znx_dft(module, self.rank() + 1, self.size());
509
510 let mut tmp: GLWECiphertextFourier<&mut [u8], FFT64> = GLWECiphertextFourier::<&mut [u8], FFT64> {
511 data: tmp_data,
512 basek: self.basek(),
513 k: self.k(),
514 };
515
516 (0..self.rank() + 1).for_each(|col_i| {
517 (0..self.rows()).for_each(|row_j| {
518 self.get_row(module, row_j, col_i, &mut tmp);
519 tmp.external_product_inplace(module, rhs, scratch1);
520 self.set_row(module, row_j, col_i, &tmp);
521 });
522 });
523 }
524}
525
526impl<DataSelf> GetRow<FFT64> for GGSWCiphertext<DataSelf, FFT64>
527where
528 MatZnxDft<DataSelf, FFT64>: MatZnxDftToRef<FFT64>,
529{
530 fn get_row<R>(&self, module: &Module<FFT64>, row_i: usize, col_j: usize, res: &mut R)
531 where
532 R: VecZnxDftToMut<FFT64>,
533 {
534 module.vmp_extract_row(res, self, row_i, col_j);
535 }
536}
537
538impl<DataSelf> SetRow<FFT64> for GGSWCiphertext<DataSelf, FFT64>
539where
540 MatZnxDft<DataSelf, FFT64>: MatZnxDftToMut<FFT64>,
541{
542 fn set_row<R>(&mut self, module: &Module<FFT64>, row_i: usize, col_j: usize, a: &R)
543 where
544 R: VecZnxDftToRef<FFT64>,
545 {
546 module.vmp_prepare_row(self, row_i, col_j, a);
547 }
548}