poulpy/core/src/glwe/encryption.rs
core/src/glwe/encryption.rsRaw
1use backend::hal::{
2 api::{
3 ScratchAvailable, SvpApply, SvpApplyInplace, SvpPPolAllocBytes, SvpPrepare, TakeScalarZnx, TakeSvpPPol, TakeVecZnx,
4 TakeVecZnxDft, VecZnxAddInplace, VecZnxAddNormal, VecZnxBigAddNormal, VecZnxBigAddSmallInplace, VecZnxBigAllocBytes,
5 VecZnxBigNormalize, VecZnxDftAllocBytes, VecZnxDftFromVecZnx, VecZnxDftToVecZnxBigConsume, VecZnxFillUniform,
6 VecZnxNormalize, VecZnxNormalizeInplace, VecZnxNormalizeTmpBytes, VecZnxSub, VecZnxSubABInplace, ZnxInfos, ZnxZero,
7 },
8 layouts::{Backend, DataMut, DataRef, Module, ScalarZnx, Scratch, VecZnx, VecZnxBig},
9};
10use sampling::source::Source;
11
12use crate::{
13 GLWECiphertext, GLWECiphertextCompressed, GLWEPlaintext, GLWEPublicKeyExec, GLWESecretExec, Infos, SIX_SIGMA,
14 dist::Distribution,
15};
16
17pub trait GLWEEncryptSkFamily<B: Backend> = VecZnxDftAllocBytes
18 + VecZnxBigNormalize<B>
19 + VecZnxDftFromVecZnx<B>
20 + SvpApplyInplace<B>
21 + VecZnxDftToVecZnxBigConsume<B>
22 + VecZnxNormalizeTmpBytes
23 + VecZnxFillUniform
24 + VecZnxSubABInplace
25 + VecZnxAddInplace
26 + VecZnxNormalizeInplace<B>
27 + VecZnxAddNormal
28 + VecZnxNormalize<B>
29 + VecZnxSub;
30
31pub trait GLWEEncryptPkFamily<B: Backend> = VecZnxDftAllocBytes
32 + VecZnxBigAllocBytes
33 + SvpPPolAllocBytes
34 + SvpPrepare<B>
35 + SvpApply<B>
36 + VecZnxDftToVecZnxBigConsume<B>
37 + VecZnxBigAddNormal<B>
38 + VecZnxBigAddSmallInplace<B>
39 + VecZnxBigNormalize<B>
40 + VecZnxNormalizeTmpBytes;
41
42impl GLWECiphertext<Vec<u8>> {
43 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
44 where
45 Module<B>: GLWEEncryptSkFamily<B>,
46 {
47 let size: usize = k.div_ceil(basek);
48 module.vec_znx_normalize_tmp_bytes(n) + 2 * VecZnx::alloc_bytes(n, 1, size) + module.vec_znx_dft_alloc_bytes(n, 1, size)
49 }
50 pub fn encrypt_pk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
51 where
52 Module<B>: GLWEEncryptPkFamily<B>,
53 {
54 let size: usize = k.div_ceil(basek);
55 ((module.vec_znx_dft_alloc_bytes(n, 1, size) + module.vec_znx_big_alloc_bytes(n, 1, size)) | ScalarZnx::alloc_bytes(n, 1))
56 + module.svp_ppol_alloc_bytes(n, 1)
57 + module.vec_znx_normalize_tmp_bytes(n)
58 }
59}
60
61impl<DataSelf: DataMut> GLWECiphertext<DataSelf> {
62 pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
63 &mut self,
64 module: &Module<B>,
65 pt: &GLWEPlaintext<DataPt>,
66 sk: &GLWESecretExec<DataSk, B>,
67 source_xa: &mut Source,
68 source_xe: &mut Source,
69 sigma: f64,
70 scratch: &mut Scratch<B>,
71 ) where
72 Module<B>: GLWEEncryptSkFamily<B>,
73 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
74 {
75 #[cfg(debug_assertions)]
76 {
77 assert_eq!(self.rank(), sk.rank());
78 assert_eq!(sk.n(), self.n());
79 assert_eq!(pt.n(), self.n());
80 assert!(
81 scratch.available() >= GLWECiphertext::encrypt_sk_scratch_space(module, self.n(), self.basek(), self.k()),
82 "scratch.available(): {} < GLWECiphertext::encrypt_sk_scratch_space: {}",
83 scratch.available(),
84 GLWECiphertext::encrypt_sk_scratch_space(module, self.n(), self.basek(), self.k())
85 )
86 }
87
88 self.encrypt_sk_internal(
89 module,
90 Some((pt, 0)),
91 sk,
92 source_xa,
93 source_xe,
94 sigma,
95 scratch,
96 );
97 }
98
99 pub fn encrypt_zero_sk<DataSk: DataRef, B: Backend>(
100 &mut self,
101 module: &Module<B>,
102 sk: &GLWESecretExec<DataSk, B>,
103 source_xa: &mut Source,
104 source_xe: &mut Source,
105 sigma: f64,
106 scratch: &mut Scratch<B>,
107 ) where
108 Module<B>: GLWEEncryptSkFamily<B>,
109 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
110 {
111 #[cfg(debug_assertions)]
112 {
113 assert_eq!(self.rank(), sk.rank());
114 assert_eq!(sk.n(), self.n());
115 assert!(
116 scratch.available() >= GLWECiphertext::encrypt_sk_scratch_space(module, self.n(), self.basek(), self.k()),
117 "scratch.available(): {} < GLWECiphertext::encrypt_sk_scratch_space: {}",
118 scratch.available(),
119 GLWECiphertext::encrypt_sk_scratch_space(module, self.n(), self.basek(), self.k())
120 )
121 }
122 self.encrypt_sk_internal(
123 module,
124 None::<(&GLWEPlaintext<Vec<u8>>, usize)>,
125 sk,
126 source_xa,
127 source_xe,
128 sigma,
129 scratch,
130 );
131 }
132
133 pub(crate) fn encrypt_sk_internal<DataPt: DataRef, DataSk: DataRef, B: Backend>(
134 &mut self,
135 module: &Module<B>,
136 pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
137 sk: &GLWESecretExec<DataSk, B>,
138 source_xa: &mut Source,
139 source_xe: &mut Source,
140 sigma: f64,
141 scratch: &mut Scratch<B>,
142 ) where
143 Module<B>: GLWEEncryptSkFamily<B>,
144 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
145 {
146 let cols: usize = self.rank() + 1;
147 encrypt_sk_internal(
148 module,
149 self.basek(),
150 self.k(),
151 &mut self.data,
152 cols,
153 false,
154 pt,
155 sk,
156 source_xa,
157 source_xe,
158 sigma,
159 scratch,
160 );
161 }
162
163 pub fn encrypt_pk<DataPt: DataRef, DataPk: DataRef, B: Backend>(
164 &mut self,
165 module: &Module<B>,
166 pt: &GLWEPlaintext<DataPt>,
167 pk: &GLWEPublicKeyExec<DataPk, B>,
168 source_xu: &mut Source,
169 source_xe: &mut Source,
170 sigma: f64,
171 scratch: &mut Scratch<B>,
172 ) where
173 Module<B>: GLWEEncryptPkFamily<B>,
174 Scratch<B>: TakeVecZnxDft<B> + TakeSvpPPol<B> + TakeScalarZnx,
175 {
176 self.encrypt_pk_internal::<DataPt, DataPk, B>(
177 module,
178 Some((pt, 0)),
179 pk,
180 source_xu,
181 source_xe,
182 sigma,
183 scratch,
184 );
185 }
186
187 pub fn encrypt_zero_pk<DataPk: DataRef, B: Backend>(
188 &mut self,
189 module: &Module<B>,
190 pk: &GLWEPublicKeyExec<DataPk, B>,
191 source_xu: &mut Source,
192 source_xe: &mut Source,
193 sigma: f64,
194 scratch: &mut Scratch<B>,
195 ) where
196 Module<B>: GLWEEncryptPkFamily<B>,
197 Scratch<B>: TakeVecZnxDft<B> + TakeSvpPPol<B> + TakeScalarZnx,
198 {
199 self.encrypt_pk_internal::<Vec<u8>, DataPk, B>(
200 module,
201 None::<(&GLWEPlaintext<Vec<u8>>, usize)>,
202 pk,
203 source_xu,
204 source_xe,
205 sigma,
206 scratch,
207 );
208 }
209
210 pub(crate) fn encrypt_pk_internal<DataPt: DataRef, DataPk: DataRef, B: Backend>(
211 &mut self,
212 module: &Module<B>,
213 pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
214 pk: &GLWEPublicKeyExec<DataPk, B>,
215 source_xu: &mut Source,
216 source_xe: &mut Source,
217 sigma: f64,
218 scratch: &mut Scratch<B>,
219 ) where
220 Module<B>: VecZnxDftAllocBytes
221 + SvpPPolAllocBytes
222 + SvpPrepare<B>
223 + SvpApply<B>
224 + VecZnxDftToVecZnxBigConsume<B>
225 + VecZnxBigAddNormal<B>
226 + VecZnxBigAddSmallInplace<B>
227 + VecZnxBigNormalize<B>,
228 Scratch<B>: TakeVecZnxDft<B> + TakeSvpPPol<B> + TakeScalarZnx,
229 {
230 #[cfg(debug_assertions)]
231 {
232 assert_eq!(self.basek(), pk.basek());
233 assert_eq!(self.n(), pk.n());
234 assert_eq!(self.rank(), pk.rank());
235 if let Some((pt, _)) = pt {
236 assert_eq!(pt.basek(), pk.basek());
237 assert_eq!(pt.n(), pk.n());
238 }
239 }
240
241 let basek: usize = pk.basek();
242 let size_pk: usize = pk.size();
243 let cols: usize = self.rank() + 1;
244
245 // Generates u according to the underlying secret distribution.
246 let (mut u_dft, scratch_1) = scratch.take_svp_ppol(self.n(), 1);
247
248 {
249 let (mut u, _) = scratch_1.take_scalar_znx(self.n(), 1);
250 match pk.dist {
251 Distribution::NONE => panic!(
252 "invalid public key: SecretDistribution::NONE, ensure it has been correctly intialized through \
253 Self::generate"
254 ),
255 Distribution::TernaryFixed(hw) => u.fill_ternary_hw(0, hw, source_xu),
256 Distribution::TernaryProb(prob) => u.fill_ternary_prob(0, prob, source_xu),
257 Distribution::BinaryFixed(hw) => u.fill_binary_hw(0, hw, source_xu),
258 Distribution::BinaryProb(prob) => u.fill_binary_prob(0, prob, source_xu),
259 Distribution::BinaryBlock(block_size) => u.fill_binary_block(0, block_size, source_xu),
260 Distribution::ZERO => {}
261 }
262
263 module.svp_prepare(&mut u_dft, 0, &u, 0);
264 }
265
266 // ct[i] = pk[i] * u + ei (+ m if col = i)
267 (0..cols).for_each(|i| {
268 let (mut ci_dft, scratch_2) = scratch_1.take_vec_znx_dft(self.n(), 1, size_pk);
269 // ci_dft = DFT(u) * DFT(pk[i])
270 module.svp_apply(&mut ci_dft, 0, &u_dft, 0, &pk.data, i);
271
272 // ci_big = u * p[i]
273 let mut ci_big = module.vec_znx_dft_to_vec_znx_big_consume(ci_dft);
274
275 // ci_big = u * pk[i] + e
276 module.vec_znx_big_add_normal(
277 basek,
278 &mut ci_big,
279 0,
280 pk.k(),
281 source_xe,
282 sigma,
283 sigma * SIX_SIGMA,
284 );
285
286 // ci_big = u * pk[i] + e + m (if col = i)
287 if let Some((pt, col)) = pt {
288 if col == i {
289 module.vec_znx_big_add_small_inplace(&mut ci_big, 0, &pt.data, 0);
290 }
291 }
292
293 // ct[i] = norm(ci_big)
294 module.vec_znx_big_normalize(basek, &mut self.data, i, &ci_big, 0, scratch_2);
295 });
296 }
297}
298
299impl GLWECiphertextCompressed<Vec<u8>> {
300 pub fn encrypt_sk_scratch_space<B: Backend>(module: &Module<B>, n: usize, basek: usize, k: usize) -> usize
301 where
302 Module<B>: GLWEEncryptSkFamily<B>,
303 {
304 GLWECiphertext::encrypt_sk_scratch_space(module, n, basek, k)
305 }
306}
307impl<D: DataMut> GLWECiphertextCompressed<D> {
308 pub fn encrypt_sk<DataPt: DataRef, DataSk: DataRef, B: Backend>(
309 &mut self,
310 module: &Module<B>,
311 pt: &GLWEPlaintext<DataPt>,
312 sk: &GLWESecretExec<DataSk, B>,
313 seed_xa: [u8; 32],
314 source_xe: &mut Source,
315 sigma: f64,
316 scratch: &mut Scratch<B>,
317 ) where
318 Module<B>: GLWEEncryptSkFamily<B>,
319 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
320 {
321 self.encrypt_sk_internal(
322 module,
323 Some((pt, 0)),
324 sk,
325 seed_xa,
326 source_xe,
327 sigma,
328 scratch,
329 );
330 }
331
332 pub(crate) fn encrypt_sk_internal<DataPt: DataRef, DataSk: DataRef, B: Backend>(
333 &mut self,
334 module: &Module<B>,
335 pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
336 sk: &GLWESecretExec<DataSk, B>,
337 seed_xa: [u8; 32],
338 source_xe: &mut Source,
339 sigma: f64,
340 scratch: &mut Scratch<B>,
341 ) where
342 Module<B>: GLWEEncryptSkFamily<B>,
343 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
344 {
345 let mut source_xa = Source::new(seed_xa);
346 let cols: usize = self.rank() + 1;
347 encrypt_sk_internal(
348 module,
349 self.basek(),
350 self.k(),
351 &mut self.data,
352 cols,
353 true,
354 pt,
355 sk,
356 &mut source_xa,
357 source_xe,
358 sigma,
359 scratch,
360 );
361 self.seed = seed_xa;
362 }
363}
364
365pub(crate) fn encrypt_sk_internal<DataCt: DataMut, DataPt: DataRef, DataSk: DataRef, B: Backend>(
366 module: &Module<B>,
367 basek: usize,
368 k: usize,
369 ct: &mut VecZnx<DataCt>,
370 cols: usize,
371 compressed: bool,
372 pt: Option<(&GLWEPlaintext<DataPt>, usize)>,
373 sk: &GLWESecretExec<DataSk, B>,
374 source_xa: &mut Source,
375 source_xe: &mut Source,
376 sigma: f64,
377 scratch: &mut Scratch<B>,
378) where
379 Module<B>: GLWEEncryptSkFamily<B>,
380 Scratch<B>: TakeVecZnxDft<B> + ScratchAvailable + TakeVecZnx,
381{
382 #[cfg(debug_assertions)]
383 {
384 if compressed {
385 assert_eq!(
386 ct.cols(),
387 1,
388 "invalid ciphertext: compressed tag=true but #cols={} != 1",
389 ct.cols()
390 )
391 }
392 }
393
394 let size: usize = ct.size();
395
396 let (mut c0, scratch_1) = scratch.take_vec_znx(ct.n(), 1, size);
397 c0.zero();
398
399 {
400 let (mut ci, scratch_2) = scratch_1.take_vec_znx(ct.n(), 1, size);
401
402 // ct[i] = uniform
403 // ct[0] -= c[i] * s[i],
404 (1..cols).for_each(|i| {
405 let col_ct: usize;
406 if compressed {
407 col_ct = 0;
408 } else {
409 col_ct = i;
410 }
411
412 // ct[i] = uniform (+ pt)
413 module.vec_znx_fill_uniform(basek, ct, col_ct, k, source_xa);
414
415 let (mut ci_dft, scratch_3) = scratch_2.take_vec_znx_dft(ct.n(), 1, size);
416
417 // ci = ct[i] - pt
418 // i.e. we act as we sample ct[i] already as uniform + pt
419 // and if there is a pt, then we subtract it before applying DFT
420 if let Some((pt, col)) = pt {
421 if i == col {
422 module.vec_znx_sub(&mut ci, 0, ct, col_ct, &pt.data, 0);
423 module.vec_znx_normalize_inplace(basek, &mut ci, 0, scratch_3);
424 module.vec_znx_dft_from_vec_znx(1, 0, &mut ci_dft, 0, &ci, 0);
425 } else {
426 module.vec_znx_dft_from_vec_znx(1, 0, &mut ci_dft, 0, ct, col_ct);
427 }
428 } else {
429 module.vec_znx_dft_from_vec_znx(1, 0, &mut ci_dft, 0, ct, col_ct);
430 }
431
432 module.svp_apply_inplace(&mut ci_dft, 0, &sk.data, i - 1);
433 let ci_big: VecZnxBig<&mut [u8], B> = module.vec_znx_dft_to_vec_znx_big_consume(ci_dft);
434
435 // use c[0] as buffer, which is overwritten later by the normalization step
436 module.vec_znx_big_normalize(basek, &mut ci, 0, &ci_big, 0, scratch_3);
437
438 // c0_tmp = -c[i] * s[i] (use c[0] as buffer)
439 module.vec_znx_sub_ab_inplace(&mut c0, 0, &ci, 0);
440 });
441 }
442
443 // c[0] += e
444 module.vec_znx_add_normal(basek, &mut c0, 0, k, source_xe, sigma, sigma * SIX_SIGMA);
445
446 // c[0] += m if col = 0
447 if let Some((pt, col)) = pt {
448 if col == 0 {
449 module.vec_znx_add_inplace(&mut c0, 0, &pt.data, 0);
450 }
451 }
452
453 // c[0] = norm(c[0])
454 module.vec_znx_normalize(basek, ct, 0, &c0, 0, scratch_1);
455}