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