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