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