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