poulpy/core/src/glwe_ciphertext.rs
core/src/glwe_ciphertext.rsRaw
1use base2k::{
2 AddNormal, Backend, FFT64, FillUniform, MatZnxDft, MatZnxDftToRef, Module, ScalarZnxAlloc, ScalarZnxDft, ScalarZnxDftAlloc,
3 ScalarZnxDftOps, ScalarZnxDftToRef, Scratch, VecZnx, VecZnxAlloc, VecZnxBig, VecZnxBigAlloc, VecZnxBigOps, VecZnxBigScratch,
4 VecZnxDft, VecZnxDftAlloc, VecZnxDftOps, VecZnxDftToMut, VecZnxDftToRef, VecZnxOps, VecZnxToMut, VecZnxToRef, ZnxInfos,
5 ZnxZero,
6};
7use sampling::source::Source;
8
9use crate::{
10 SIX_SIGMA,
11 elem::Infos,
12 gglwe_ciphertext::GGLWECiphertext,
13 ggsw_ciphertext::GGSWCiphertext,
14 glwe_ciphertext_fourier::GLWECiphertextFourier,
15 glwe_plaintext::GLWEPlaintext,
16 keys::{GLWEPublicKey, SecretDistribution, SecretKeyFourier},
17 keyswitch_key::GLWESwitchingKey,
18 utils::derive_size,
19 vec_glwe_product::{VecGLWEProduct, VecGLWEProductScratchSpace},
20};
21
22pub struct GLWECiphertext<C> {
23 pub data: VecZnx<C>,
24 pub basek: usize,
25 pub k: usize,
26}
27
28impl GLWECiphertext<Vec<u8>> {
29 pub fn new<B: Backend>(module: &Module<B>, basek: usize, k: usize, rank: usize) -> Self {
30 Self {
31 data: module.new_vec_znx(rank + 1, derive_size(basek, k)),
32 basek,
33 k,
34 }
35 }
36}
37
38impl<T> Infos for GLWECiphertext<T> {
39 type Inner = VecZnx<T>;
40
41 fn inner(&self) -> &Self::Inner {
42 &self.data
43 }
44
45 fn basek(&self) -> usize {
46 self.basek
47 }
48
49 fn k(&self) -> usize {
50 self.k
51 }
52}
53
54impl<T> GLWECiphertext<T> {
55 pub fn rank(&self) -> usize {
56 self.cols() - 1
57 }
58}
59
60impl<C> VecZnxToMut for GLWECiphertext<C>
61where
62 VecZnx<C>: VecZnxToMut,
63{
64 fn to_mut(&mut self) -> VecZnx<&mut [u8]> {
65 self.data.to_mut()
66 }
67}
68
69impl<C> VecZnxToRef for GLWECiphertext<C>
70where
71 VecZnx<C>: VecZnxToRef,
72{
73 fn to_ref(&self) -> VecZnx<&[u8]> {
74 self.data.to_ref()
75 }
76}
77
78impl<C> GLWECiphertext<C>
79where
80 VecZnx<C>: VecZnxToRef,
81{
82 #[allow(dead_code)]
83 pub(crate) fn dft<R>(&self, module: &Module<FFT64>, res: &mut GLWECiphertextFourier<R, FFT64>)
84 where
85 VecZnxDft<R, FFT64>: VecZnxDftToMut<FFT64> + ZnxInfos,
86 {
87 #[cfg(debug_assertions)]
88 {
89 assert_eq!(self.rank(), res.rank());
90 assert_eq!(self.basek(), res.basek())
91 }
92
93 (0..self.rank() + 1).for_each(|i| {
94 module.vec_znx_dft(res, i, self, i);
95 })
96 }
97}
98
99impl GLWECiphertext<Vec<u8>> {
100 pub fn encrypt_sk_scratch_space(module: &Module<FFT64>, _rank: usize, ct_size: usize) -> usize {
101 module.vec_znx_big_normalize_tmp_bytes()
102 + module.bytes_of_vec_znx_dft(1, ct_size)
103 + module.bytes_of_vec_znx_big(1, ct_size)
104 }
105 pub fn encrypt_pk_scratch_space(module: &Module<FFT64>, _rank: usize, pk_size: usize) -> usize {
106 ((module.bytes_of_vec_znx_dft(1, pk_size) + module.bytes_of_vec_znx_big(1, pk_size)) | module.bytes_of_scalar_znx(1))
107 + module.bytes_of_scalar_znx_dft(1)
108 + module.vec_znx_big_normalize_tmp_bytes()
109 }
110
111 pub fn decrypt_scratch_space(module: &Module<FFT64>, ct_size: usize) -> usize {
112 (module.vec_znx_big_normalize_tmp_bytes() | module.bytes_of_vec_znx_dft(1, ct_size))
113 + module.bytes_of_vec_znx_big(1, ct_size)
114 }
115
116 pub fn keyswitch_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
117 <GGLWECiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_glwe_scratch_space(module, res_size, lhs, rhs)
118 }
119
120 pub fn keyswitch_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize) -> usize {
121 <GGLWECiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_glwe_inplace_scratch_space(
122 module, res_size, rhs,
123 )
124 }
125
126 pub fn external_product_scratch_space(module: &Module<FFT64>, res_size: usize, lhs: usize, rhs: usize) -> usize {
127 <GGSWCiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_glwe_scratch_space(module, res_size, lhs, rhs)
128 }
129
130 pub fn external_product_inplace_scratch_space(module: &Module<FFT64>, res_size: usize, rhs: usize) -> usize {
131 <GGSWCiphertext<Vec<u8>, FFT64> as VecGLWEProductScratchSpace>::prod_with_glwe_inplace_scratch_space(
132 module, res_size, rhs,
133 )
134 }
135}
136
137impl<DataSelf> GLWECiphertext<DataSelf>
138where
139 VecZnx<DataSelf>: VecZnxToMut + VecZnxToRef,
140{
141 pub fn encrypt_sk<DataPt, DataSk>(
142 &mut self,
143 module: &Module<FFT64>,
144 pt: &GLWEPlaintext<DataPt>,
145 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
146 source_xa: &mut Source,
147 source_xe: &mut Source,
148 sigma: f64,
149 scratch: &mut Scratch,
150 ) where
151 VecZnx<DataPt>: VecZnxToRef,
152 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
153 {
154 self.encrypt_sk_private(
155 module,
156 Some((pt, 0)),
157 sk_dft,
158 source_xa,
159 source_xe,
160 sigma,
161 scratch,
162 );
163 }
164
165 pub fn encrypt_zero_sk<DataSk>(
166 &mut self,
167 module: &Module<FFT64>,
168 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
169 source_xa: &mut Source,
170 source_xe: &mut Source,
171 sigma: f64,
172 scratch: &mut Scratch,
173 ) where
174 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
175 {
176 self.encrypt_sk_private(module, None, sk_dft, source_xa, source_xe, sigma, scratch);
177 }
178
179 pub fn encrypt_pk<DataPt, DataPk>(
180 &mut self,
181 module: &Module<FFT64>,
182 pt: &GLWEPlaintext<DataPt>,
183 pk: &GLWEPublicKey<DataPk, FFT64>,
184 source_xu: &mut Source,
185 source_xe: &mut Source,
186 sigma: f64,
187 scratch: &mut Scratch,
188 ) where
189 VecZnx<DataPt>: VecZnxToRef,
190 VecZnxDft<DataPk, FFT64>: VecZnxDftToRef<FFT64>,
191 {
192 self.encrypt_pk_private(
193 module,
194 Some((pt, 0)),
195 pk,
196 source_xu,
197 source_xe,
198 sigma,
199 scratch,
200 );
201 }
202
203 pub fn encrypt_zero_pk<DataPk>(
204 &mut self,
205 module: &Module<FFT64>,
206 pk: &GLWEPublicKey<DataPk, FFT64>,
207 source_xu: &mut Source,
208 source_xe: &mut Source,
209 sigma: f64,
210 scratch: &mut Scratch,
211 ) where
212 VecZnxDft<DataPk, FFT64>: VecZnxDftToRef<FFT64>,
213 {
214 self.encrypt_pk_private(module, None, pk, source_xu, source_xe, sigma, scratch);
215 }
216
217 pub fn keyswitch<DataLhs, DataRhs>(
218 &mut self,
219 module: &Module<FFT64>,
220 lhs: &GLWECiphertext<DataLhs>,
221 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
222 scratch: &mut Scratch,
223 ) where
224 VecZnx<DataLhs>: VecZnxToRef,
225 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
226 {
227 rhs.0.prod_with_glwe(module, self, lhs, scratch);
228 }
229
230 pub fn keyswitch_inplace<DataRhs>(
231 &mut self,
232 module: &Module<FFT64>,
233 rhs: &GLWESwitchingKey<DataRhs, FFT64>,
234 scratch: &mut Scratch,
235 ) where
236 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
237 {
238 rhs.0.prod_with_glwe_inplace(module, self, scratch);
239 }
240
241 pub fn external_product<DataLhs, DataRhs>(
242 &mut self,
243 module: &Module<FFT64>,
244 lhs: &GLWECiphertext<DataLhs>,
245 rhs: &GGSWCiphertext<DataRhs, FFT64>,
246 scratch: &mut Scratch,
247 ) where
248 VecZnx<DataLhs>: VecZnxToRef,
249 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
250 {
251 rhs.prod_with_glwe(module, self, lhs, scratch);
252 }
253
254 pub fn external_product_inplace<DataRhs>(
255 &mut self,
256 module: &Module<FFT64>,
257 rhs: &GGSWCiphertext<DataRhs, FFT64>,
258 scratch: &mut Scratch,
259 ) where
260 MatZnxDft<DataRhs, FFT64>: MatZnxDftToRef<FFT64>,
261 {
262 rhs.prod_with_glwe_inplace(module, self, scratch);
263 }
264
265 pub(crate) fn encrypt_sk_private<DataPt, DataSk>(
266 &mut self,
267 module: &Module<FFT64>,
268 pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
269 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
270 source_xa: &mut Source,
271 source_xe: &mut Source,
272 sigma: f64,
273 scratch: &mut Scratch,
274 ) where
275 VecZnx<DataPt>: VecZnxToRef,
276 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
277 {
278 #[cfg(debug_assertions)]
279 {
280 assert_eq!(self.rank(), sk_dft.rank());
281 assert_eq!(sk_dft.n(), module.n());
282 assert_eq!(self.n(), module.n());
283 if let Some((pt, col)) = pt {
284 assert_eq!(pt.n(), module.n());
285 assert!(col < self.rank() + 1);
286 }
287 }
288
289 let log_base2k: usize = self.basek();
290 let log_k: usize = self.k();
291 let size: usize = self.size();
292 let cols: usize = self.rank() + 1;
293
294 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx(module, 1, size);
295 c0_big.zero();
296
297 {
298 // c[i] = uniform
299 // c[0] -= c[i] * s[i],
300 (1..cols).for_each(|i| {
301 let (mut ci_dft, scratch_2) = scratch_1.tmp_vec_znx_dft(module, 1, size);
302
303 // c[i] = uniform
304 self.data.fill_uniform(log_base2k, i, size, source_xa);
305
306 // c[i] = norm(IDFT(DFT(c[i]) * DFT(s[i])))
307 module.vec_znx_dft(&mut ci_dft, 0, self, i);
308 module.svp_apply_inplace(&mut ci_dft, 0, sk_dft, i - 1);
309 let ci_big: VecZnxBig<&mut [u8], FFT64> = module.vec_znx_idft_consume(ci_dft);
310
311 // use c[0] as buffer, which is overwritten later by the normalization step
312 module.vec_znx_big_normalize(log_base2k, self, 0, &ci_big, 0, scratch_2);
313
314 // c0_tmp = -c[i] * s[i] (use c[0] as buffer)
315 module.vec_znx_sub_ab_inplace(&mut c0_big, 0, self, 0);
316
317 // c[i] += m if col = i
318 if let Some((pt, col)) = pt {
319 if i == col {
320 module.vec_znx_add_inplace(self, i, pt, 0);
321 module.vec_znx_normalize_inplace(log_base2k, self, i, scratch_2);
322 }
323 }
324 });
325 }
326
327 // c[0] += e
328 c0_big.add_normal(log_base2k, 0, log_k, source_xe, sigma, sigma * SIX_SIGMA);
329
330 // c[0] += m if col = 0
331 if let Some((pt, col)) = pt {
332 if col == 0 {
333 module.vec_znx_add_inplace(&mut c0_big, 0, pt, 0);
334 }
335 }
336
337 // c[0] = norm(c[0])
338 module.vec_znx_normalize(log_base2k, self, 0, &c0_big, 0, scratch_1);
339 }
340
341 pub(crate) fn encrypt_pk_private<DataPt, DataPk>(
342 &mut self,
343 module: &Module<FFT64>,
344 pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
345 pk: &GLWEPublicKey<DataPk, FFT64>,
346 source_xu: &mut Source,
347 source_xe: &mut Source,
348 sigma: f64,
349 scratch: &mut Scratch,
350 ) where
351 VecZnx<DataPt>: VecZnxToRef,
352 VecZnxDft<DataPk, FFT64>: VecZnxDftToRef<FFT64>,
353 {
354 #[cfg(debug_assertions)]
355 {
356 assert_eq!(self.basek(), pk.basek());
357 assert_eq!(self.n(), module.n());
358 assert_eq!(pk.n(), module.n());
359 assert_eq!(self.rank(), pk.rank());
360 if let Some((pt, _)) = pt {
361 assert_eq!(pt.basek(), pk.basek());
362 assert_eq!(pt.n(), module.n());
363 }
364 }
365
366 let log_base2k: usize = pk.basek();
367 let size_pk: usize = pk.size();
368 let cols: usize = self.rank() + 1;
369
370 // Generates u according to the underlying secret distribution.
371 let (mut u_dft, scratch_1) = scratch.tmp_scalar_znx_dft(module, 1);
372
373 {
374 let (mut u, _) = scratch_1.tmp_scalar_znx(module, 1);
375 match pk.dist {
376 SecretDistribution::NONE => panic!(
377 "invalid public key: SecretDistribution::NONE, ensure it has been correctly intialized through \
378 Self::generate"
379 ),
380 SecretDistribution::TernaryFixed(hw) => u.fill_ternary_hw(0, hw, source_xu),
381 SecretDistribution::TernaryProb(prob) => u.fill_ternary_prob(0, prob, source_xu),
382 SecretDistribution::ZERO => {}
383 }
384
385 module.svp_prepare(&mut u_dft, 0, &u, 0);
386 }
387
388 // ct[i] = pk[i] * u + ei (+ m if col = i)
389 (0..cols).for_each(|i| {
390 let (mut ci_dft, scratch_2) = scratch_1.tmp_vec_znx_dft(module, 1, size_pk);
391 // ci_dft = DFT(u) * DFT(pk[i])
392 module.svp_apply(&mut ci_dft, 0, &u_dft, 0, pk, i);
393
394 // ci_big = u * p[i]
395 let mut ci_big = module.vec_znx_idft_consume(ci_dft);
396
397 // ci_big = u * pk[i] + e
398 ci_big.add_normal(log_base2k, 0, pk.k(), source_xe, sigma, sigma * SIX_SIGMA);
399
400 // ci_big = u * pk[i] + e + m (if col = i)
401 if let Some((pt, col)) = pt {
402 if col == i {
403 module.vec_znx_big_add_small_inplace(&mut ci_big, 0, pt, 0);
404 }
405 }
406
407 // ct[i] = norm(ci_big)
408 module.vec_znx_big_normalize(log_base2k, self, i, &ci_big, 0, scratch_2);
409 });
410 }
411}
412
413impl<DataSelf> GLWECiphertext<DataSelf>
414where
415 VecZnx<DataSelf>: VecZnxToRef,
416{
417 pub fn decrypt<DataPt, DataSk>(
418 &self,
419 module: &Module<FFT64>,
420 pt: &mut GLWEPlaintext<DataPt>,
421 sk_dft: &SecretKeyFourier<DataSk, FFT64>,
422 scratch: &mut Scratch,
423 ) where
424 VecZnx<DataPt>: VecZnxToMut,
425 ScalarZnxDft<DataSk, FFT64>: ScalarZnxDftToRef<FFT64>,
426 {
427 #[cfg(debug_assertions)]
428 {
429 assert_eq!(self.rank(), sk_dft.rank());
430 assert_eq!(self.n(), module.n());
431 assert_eq!(pt.n(), module.n());
432 assert_eq!(sk_dft.n(), module.n());
433 }
434
435 let cols: usize = self.rank() + 1;
436
437 let (mut c0_big, scratch_1) = scratch.tmp_vec_znx_big(module, 1, self.size()); // TODO optimize size when pt << ct
438 c0_big.zero();
439
440 {
441 (1..cols).for_each(|i| {
442 // ci_dft = DFT(a[i]) * DFT(s[i])
443 let (mut ci_dft, _) = scratch_1.tmp_vec_znx_dft(module, 1, self.size()); // TODO optimize size when pt << ct
444 module.vec_znx_dft(&mut ci_dft, 0, self, i);
445 module.svp_apply_inplace(&mut ci_dft, 0, sk_dft, i - 1);
446 let ci_big = module.vec_znx_idft_consume(ci_dft);
447
448 // c0_big += a[i] * s[i]
449 module.vec_znx_big_add_inplace(&mut c0_big, 0, &ci_big, 0);
450 });
451 }
452
453 // c0_big = (a * s) + (-a * s + m + e) = BIG(m + e)
454 module.vec_znx_big_add_small_inplace(&mut c0_big, 0, self, 0);
455
456 // pt = norm(BIG(m + e))
457 module.vec_znx_big_normalize(self.basek(), pt, 0, &mut c0_big, 0, scratch_1);
458
459 pt.basek = self.basek();
460 pt.k = pt.k().min(self.k());
461 }
462}