phantom-zone/src/bool/keys.rs
src/bool/keys.rsRaw
1use std::{collections::HashMap, hash::Hash, marker::PhantomData};
2
3use crate::{
4 backend::{ModInit, VectorOps},
5 lwe::LweSecret,
6 random::{NewWithSeed, RandomFillUniformInModulus},
7 rgsw::RlweSecret,
8 utils::WithLocal,
9 Decryptor, Encryptor, Matrix, MatrixEntity, MatrixMut, MultiPartyDecryptor, RowEntity, RowMut,
10};
11
12use super::{parameters, BoolEvaluator, BoolParameters, CiphertextModulus};
13
14/// Client key with RLWE and LWE secrets
15#[derive(Clone)]
16pub struct ClientKey {
17 sk_rlwe: RlweSecret,
18 sk_lwe: LweSecret,
19}
20
21mod impl_ck {
22 use super::*;
23
24 // Client key
25 impl ClientKey {
26 pub(in super::super) fn random() -> Self {
27 let sk_rlwe = RlweSecret::random(0, 0);
28 let sk_lwe = LweSecret::random(0, 0);
29 Self { sk_rlwe, sk_lwe }
30 }
31
32 pub(in super::super) fn new(sk_rlwe: RlweSecret, sk_lwe: LweSecret) -> Self {
33 Self { sk_rlwe, sk_lwe }
34 }
35
36 pub(in super::super) fn sk_rlwe(&self) -> &RlweSecret {
37 &self.sk_rlwe
38 }
39
40 pub(in super::super) fn sk_lwe(&self) -> &LweSecret {
41 &self.sk_lwe
42 }
43 }
44
45 impl Encryptor<bool, Vec<u64>> for ClientKey {
46 fn encrypt(&self, m: &bool) -> Vec<u64> {
47 BoolEvaluator::with_local(|e| e.sk_encrypt(*m, self))
48 }
49 }
50
51 impl Decryptor<bool, Vec<u64>> for ClientKey {
52 fn decrypt(&self, c: &Vec<u64>) -> bool {
53 BoolEvaluator::with_local(|e| e.sk_decrypt(c, self))
54 }
55 }
56
57 impl MultiPartyDecryptor<bool, Vec<u64>> for ClientKey {
58 type DecryptionShare = u64;
59
60 fn gen_decryption_share(&self, c: &Vec<u64>) -> Self::DecryptionShare {
61 BoolEvaluator::with_local(|e| e.multi_party_decryption_share(c, &self))
62 }
63
64 fn aggregate_decryption_shares(
65 &self,
66 c: &Vec<u64>,
67 shares: &[Self::DecryptionShare],
68 ) -> bool {
69 BoolEvaluator::with_local(|e| e.multi_party_decrypt(shares, c))
70 }
71 }
72}
73
74/// Public key
75pub struct PublicKey<M, Rng, ModOp> {
76 key: M,
77 _phantom: PhantomData<(Rng, ModOp)>,
78}
79
80pub(super) mod impl_pk {
81 use super::*;
82
83 impl<M, R, Mo> PublicKey<M, R, Mo> {
84 pub(in super::super) fn key(&self) -> &M {
85 &self.key
86 }
87 }
88
89 impl<Rng, ModOp> Encryptor<bool, Vec<u64>> for PublicKey<Vec<Vec<u64>>, Rng, ModOp> {
90 fn encrypt(&self, m: &bool) -> Vec<u64> {
91 BoolEvaluator::with_local(|e| e.pk_encrypt(&self.key, *m))
92 }
93 }
94
95 impl<Rng, ModOp> Encryptor<[bool], Vec<Vec<u64>>> for PublicKey<Vec<Vec<u64>>, Rng, ModOp> {
96 fn encrypt(&self, m: &[bool]) -> Vec<Vec<u64>> {
97 BoolEvaluator::with_local(|e| e.pk_encrypt_batched(&self.key, m))
98 }
99 }
100
101 impl<
102 M: MatrixMut + MatrixEntity,
103 Rng: NewWithSeed
104 + RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>,
105 ModOp,
106 > From<SeededPublicKey<M::R, Rng::Seed, BoolParameters<M::MatElement>, ModOp>>
107 for PublicKey<M, Rng, ModOp>
108 where
109 <M as Matrix>::R: RowMut,
110 M::MatElement: Copy,
111 {
112 fn from(
113 value: SeededPublicKey<M::R, Rng::Seed, BoolParameters<M::MatElement>, ModOp>,
114 ) -> Self {
115 let mut prng = Rng::new_with_seed(value.seed);
116
117 let mut key = M::zeros(2, value.parameters.rlwe_n().0);
118 // sample A
119 RandomFillUniformInModulus::random_fill(
120 &mut prng,
121 value.parameters.rlwe_q(),
122 key.get_row_mut(0),
123 );
124 // Copy over B
125 key.get_row_mut(1).copy_from_slice(value.part_b.as_ref());
126
127 PublicKey {
128 key,
129 _phantom: PhantomData,
130 }
131 }
132 }
133
134 impl<
135 M: MatrixMut + MatrixEntity,
136 Rng: NewWithSeed
137 + RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>,
138 ModOp: VectorOps<Element = M::MatElement> + ModInit<M = CiphertextModulus<M::MatElement>>,
139 >
140 From<
141 &[CommonReferenceSeededCollectivePublicKeyShare<
142 M::R,
143 Rng::Seed,
144 BoolParameters<M::MatElement>,
145 >],
146 > for PublicKey<M, Rng, ModOp>
147 where
148 <M as Matrix>::R: RowMut,
149 Rng::Seed: Copy + PartialEq,
150 M::MatElement: PartialEq + Copy,
151 {
152 fn from(
153 value: &[CommonReferenceSeededCollectivePublicKeyShare<
154 M::R,
155 Rng::Seed,
156 BoolParameters<M::MatElement>,
157 >],
158 ) -> Self {
159 assert!(value.len() > 0);
160
161 let parameters = &value[0].parameters;
162 let mut key = M::zeros(2, parameters.rlwe_n().0);
163
164 // sample A
165 let seed = value[0].cr_seed;
166 let mut main_rng = Rng::new_with_seed(seed);
167 RandomFillUniformInModulus::random_fill(
168 &mut main_rng,
169 parameters.rlwe_q(),
170 key.get_row_mut(0),
171 );
172
173 // Sum all Bs
174 let rlweq_modop = ModOp::new(parameters.rlwe_q().clone());
175 value.iter().for_each(|share_i| {
176 assert!(share_i.cr_seed == seed);
177 assert!(&share_i.parameters == parameters);
178
179 rlweq_modop.elwise_add_mut(key.get_row_mut(1), share_i.share.as_ref());
180 });
181
182 PublicKey {
183 key,
184 _phantom: PhantomData,
185 }
186 }
187 }
188}
189
190/// Seeded public key
191struct SeededPublicKey<Ro, S, P, ModOp> {
192 part_b: Ro,
193 seed: S,
194 parameters: P,
195 _phantom: PhantomData<ModOp>,
196}
197
198mod impl_seeded_pk {
199 use super::*;
200
201 impl<R, S, ModOp>
202 From<&[CommonReferenceSeededCollectivePublicKeyShare<R, S, BoolParameters<R::Element>>]>
203 for SeededPublicKey<R, S, BoolParameters<R::Element>, ModOp>
204 where
205 ModOp: VectorOps<Element = R::Element> + ModInit<M = CiphertextModulus<R::Element>>,
206 S: PartialEq + Clone,
207 R: RowMut + RowEntity + Clone,
208 R::Element: Clone + PartialEq,
209 {
210 fn from(
211 value: &[CommonReferenceSeededCollectivePublicKeyShare<
212 R,
213 S,
214 BoolParameters<R::Element>,
215 >],
216 ) -> Self {
217 assert!(value.len() > 0);
218
219 let parameters = &value[0].parameters;
220 let cr_seed = value[0].cr_seed.clone();
221
222 // Sum all Bs
223 let rlweq_modop = ModOp::new(parameters.rlwe_q().clone());
224 let mut part_b = value[0].share.clone();
225 value.iter().skip(1).for_each(|share_i| {
226 assert!(&share_i.cr_seed == &cr_seed);
227 assert!(&share_i.parameters == parameters);
228
229 rlweq_modop.elwise_add_mut(part_b.as_mut(), share_i.share.as_ref());
230 });
231
232 Self {
233 part_b,
234 seed: cr_seed,
235 parameters: parameters.clone(),
236 _phantom: PhantomData,
237 }
238 }
239 }
240}
241
242/// CRS seeded collective public key share
243pub struct CommonReferenceSeededCollectivePublicKeyShare<Ro, S, P> {
244 share: Ro,
245 cr_seed: S,
246 parameters: P,
247}
248impl<Ro, S, P> CommonReferenceSeededCollectivePublicKeyShare<Ro, S, P> {
249 pub(super) fn new(share: Ro, cr_seed: S, parameters: P) -> Self {
250 CommonReferenceSeededCollectivePublicKeyShare {
251 share,
252 cr_seed,
253 parameters,
254 }
255 }
256}
257
258/// CRS seeded Multi-party server key share
259pub struct CommonReferenceSeededMultiPartyServerKeyShare<M: Matrix, P, S> {
260 rgsw_cts: Vec<M>,
261 /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
262 /// to -g is at 0
263 auto_keys: HashMap<usize, M>,
264 lwe_ksk: M::R,
265 /// Common reference seed
266 cr_seed: S,
267 parameters: P,
268}
269
270impl<M: Matrix, P, S> CommonReferenceSeededMultiPartyServerKeyShare<M, P, S> {
271 pub(super) fn new(
272 rgsw_cts: Vec<M>,
273 auto_keys: HashMap<usize, M>,
274 lwe_ksk: M::R,
275 cr_seed: S,
276 parameters: P,
277 ) -> Self {
278 CommonReferenceSeededMultiPartyServerKeyShare {
279 rgsw_cts,
280 auto_keys,
281 lwe_ksk,
282 cr_seed,
283 parameters,
284 }
285 }
286
287 pub(super) fn cr_seed(&self) -> &S {
288 &self.cr_seed
289 }
290
291 pub(super) fn parameters(&self) -> &P {
292 &self.parameters
293 }
294
295 pub(super) fn auto_keys(&self) -> &HashMap<usize, M> {
296 &self.auto_keys
297 }
298
299 pub(super) fn rgsw_cts(&self) -> &[M] {
300 &self.rgsw_cts
301 }
302
303 pub(super) fn lwe_ksk(&self) -> &M::R {
304 &self.lwe_ksk
305 }
306}
307
308/// CRS seeded MultiParty server key
309pub struct SeededMultiPartyServerKey<M: Matrix, S, P> {
310 rgsw_cts: Vec<M>,
311 /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
312 /// to -g is at 0
313 auto_keys: HashMap<usize, M>,
314 lwe_ksk: M::R,
315 cr_seed: S,
316 parameters: P,
317}
318
319impl<M: Matrix, S, P> SeededMultiPartyServerKey<M, S, P> {
320 pub(super) fn new(
321 rgsw_cts: Vec<M>,
322 auto_keys: HashMap<usize, M>,
323 lwe_ksk: M::R,
324 cr_seed: S,
325 parameters: P,
326 ) -> Self {
327 SeededMultiPartyServerKey {
328 rgsw_cts,
329 auto_keys,
330 lwe_ksk,
331 cr_seed,
332 parameters,
333 }
334 }
335
336 pub(super) fn rgsw_cts(&self) -> &[M] {
337 &self.rgsw_cts
338 }
339}
340
341/// Seeded single party server key
342pub struct SeededServerKey<M: Matrix, P, S> {
343 /// Rgsw cts of LWE secret elements
344 pub(crate) rgsw_cts: Vec<M>,
345 /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
346 /// to -g is at 0
347 pub(crate) auto_keys: HashMap<usize, M>,
348 /// LWE ksk to key switching LWE ciphertext from RLWE secret to LWE secret
349 pub(crate) lwe_ksk: M::R,
350 /// Parameters
351 pub(crate) parameters: P,
352 /// Main seed
353 pub(crate) seed: S,
354}
355impl<M: Matrix, S> SeededServerKey<M, BoolParameters<M::MatElement>, S> {
356 pub(super) fn from_raw(
357 auto_keys: HashMap<usize, M>,
358 rgsw_cts: Vec<M>,
359 lwe_ksk: M::R,
360 parameters: BoolParameters<M::MatElement>,
361 seed: S,
362 ) -> Self {
363 // sanity checks
364 auto_keys.iter().for_each(|v| {
365 assert!(
366 v.1.dimension()
367 == (
368 parameters.auto_decomposition_count().0,
369 parameters.rlwe_n().0
370 )
371 )
372 });
373
374 let (part_a_d, part_b_d) = parameters.rlwe_rgsw_decomposition_count();
375 rgsw_cts.iter().for_each(|v| {
376 assert!(v.dimension() == (part_a_d.0 * 2 + part_b_d.0, parameters.rlwe_n().0))
377 });
378 assert!(
379 lwe_ksk.as_ref().len()
380 == (parameters.lwe_decomposition_count().0 * parameters.rlwe_n().0)
381 );
382
383 SeededServerKey {
384 rgsw_cts,
385 auto_keys,
386 lwe_ksk,
387 parameters,
388 seed,
389 }
390 }
391}
392
393/// Server key in evaluation domain
394pub(crate) struct ServerKeyEvaluationDomain<M, R, N> {
395 /// Rgsw cts of LWE secret elements
396 rgsw_cts: Vec<M>,
397 /// Auto keys. Key corresponding to g^{k} is at index `k`. Key corresponding
398 /// to -g is at 0
399 galois_keys: HashMap<usize, M>,
400 /// LWE ksk to key switching LWE ciphertext from RLWE secret to LWE secret
401 lwe_ksk: M,
402 _phanton: PhantomData<(R, N)>,
403}
404
405pub(super) mod impl_server_key_eval_domain {
406 use itertools::{izip, Itertools};
407
408 use crate::{
409 ntt::{Ntt, NttInit},
410 pbs::PbsKey,
411 };
412
413 use super::*;
414
415 impl<M, R, N> ServerKeyEvaluationDomain<M, R, N> {
416 pub(in super::super) fn rgsw_cts(&self) -> &[M] {
417 &self.rgsw_cts
418 }
419 }
420
421 impl<
422 M: MatrixMut + MatrixEntity,
423 R: RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>
424 + NewWithSeed,
425 N: NttInit<CiphertextModulus<M::MatElement>> + Ntt<Element = M::MatElement>,
426 > From<&SeededServerKey<M, BoolParameters<M::MatElement>, R::Seed>>
427 for ServerKeyEvaluationDomain<M, R, N>
428 where
429 <M as Matrix>::R: RowMut,
430 M::MatElement: Copy,
431 R::Seed: Clone,
432 {
433 fn from(value: &SeededServerKey<M, BoolParameters<M::MatElement>, R::Seed>) -> Self {
434 let mut main_prng = R::new_with_seed(value.seed.clone());
435 let parameters = &value.parameters;
436 let g = parameters.g() as isize;
437 let ring_size = value.parameters.rlwe_n().0;
438 let lwe_n = value.parameters.lwe_n().0;
439 let rlwe_q = value.parameters.rlwe_q();
440 let lwq_q = value.parameters.lwe_q();
441
442 let nttop = N::new(rlwe_q, ring_size);
443
444 // galois keys
445 let mut auto_keys = HashMap::new();
446 let auto_decomp_count = parameters.auto_decomposition_count().0;
447 let auto_element_dlogs = parameters.auto_element_dlogs();
448 for i in auto_element_dlogs.into_iter() {
449 let seeded_auto_key = value.auto_keys.get(&i).unwrap();
450 assert!(seeded_auto_key.dimension() == (auto_decomp_count, ring_size));
451
452 let mut data = M::zeros(auto_decomp_count * 2, ring_size);
453
454 // sample RLWE'_A(-s(X^k))
455 data.iter_rows_mut().take(auto_decomp_count).for_each(|ri| {
456 RandomFillUniformInModulus::random_fill(&mut main_prng, &rlwe_q, ri.as_mut())
457 });
458
459 // copy over RLWE'B_(-s(X^k))
460 izip!(
461 data.iter_rows_mut().skip(auto_decomp_count),
462 seeded_auto_key.iter_rows()
463 )
464 .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
465
466 // Send to Evaluation domain
467 data.iter_rows_mut()
468 .for_each(|ri| nttop.forward(ri.as_mut()));
469
470 auto_keys.insert(i, data);
471 }
472
473 // RGSW ciphertexts
474 let (rlrg_a_decomp, rlrg_b_decomp) = parameters.rlwe_rgsw_decomposition_count();
475 let rgsw_cts = value
476 .rgsw_cts
477 .iter()
478 .map(|seeded_rgsw_si| {
479 assert!(
480 seeded_rgsw_si.dimension()
481 == (rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0, ring_size)
482 );
483
484 let mut data = M::zeros(rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0 * 2, ring_size);
485
486 // copy over RLWE'(-sm)
487 izip!(
488 data.iter_rows_mut().take(rlrg_a_decomp.0 * 2),
489 seeded_rgsw_si.iter_rows().take(rlrg_a_decomp.0 * 2)
490 )
491 .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
492
493 // sample RLWE'_A(m)
494 data.iter_rows_mut()
495 .skip(rlrg_a_decomp.0 * 2)
496 .take(rlrg_b_decomp.0)
497 .for_each(|ri| {
498 RandomFillUniformInModulus::random_fill(
499 &mut main_prng,
500 &rlwe_q,
501 ri.as_mut(),
502 )
503 });
504
505 // copy over RLWE'_B(m)
506 izip!(
507 data.iter_rows_mut()
508 .skip(rlrg_a_decomp.0 * 2 + rlrg_b_decomp.0),
509 seeded_rgsw_si.iter_rows().skip(rlrg_a_decomp.0 * 2)
510 )
511 .for_each(|(to_ri, from_ri)| to_ri.as_mut().copy_from_slice(from_ri.as_ref()));
512
513 // send polynomials to evaluation domain
514 data.iter_rows_mut()
515 .for_each(|ri| nttop.forward(ri.as_mut()));
516
517 data
518 })
519 .collect_vec();
520
521 // LWE ksk
522 let lwe_ksk = {
523 let d = parameters.lwe_decomposition_count().0;
524 assert!(value.lwe_ksk.as_ref().len() == d * ring_size);
525
526 let mut data = M::zeros(d * ring_size, lwe_n + 1);
527 izip!(data.iter_rows_mut(), value.lwe_ksk.as_ref().iter()).for_each(
528 |(lwe_i, bi)| {
529 RandomFillUniformInModulus::random_fill(
530 &mut main_prng,
531 &lwq_q,
532 &mut lwe_i.as_mut()[1..],
533 );
534 lwe_i.as_mut()[0] = *bi;
535 },
536 );
537
538 data
539 };
540
541 ServerKeyEvaluationDomain {
542 rgsw_cts,
543 galois_keys: auto_keys,
544 lwe_ksk,
545 _phanton: PhantomData,
546 }
547 }
548 }
549
550 impl<
551 M: MatrixMut + MatrixEntity,
552 Rng: NewWithSeed,
553 N: NttInit<CiphertextModulus<M::MatElement>> + Ntt<Element = M::MatElement>,
554 > From<&SeededMultiPartyServerKey<M, Rng::Seed, BoolParameters<M::MatElement>>>
555 for ServerKeyEvaluationDomain<M, Rng, N>
556 where
557 <M as Matrix>::R: RowMut,
558 Rng::Seed: Copy,
559 Rng: RandomFillUniformInModulus<[M::MatElement], CiphertextModulus<M::MatElement>>,
560 M::MatElement: Copy,
561 {
562 fn from(
563 value: &SeededMultiPartyServerKey<M, Rng::Seed, BoolParameters<M::MatElement>>,
564 ) -> Self {
565 let g = value.parameters.g() as isize;
566 let rlwe_n = value.parameters.rlwe_n().0;
567 let lwe_n = value.parameters.lwe_n().0;
568 let rlwe_q = value.parameters.rlwe_q();
569 let lwe_q = value.parameters.lwe_q();
570
571 let mut main_prng = Rng::new_with_seed(value.cr_seed);
572
573 let rlwe_nttop = N::new(rlwe_q, rlwe_n);
574
575 // auto keys
576 let mut auto_keys = HashMap::new();
577 let auto_d_count = value.parameters.auto_decomposition_count().0;
578 let auto_element_dlogs = value.parameters.auto_element_dlogs();
579 for i in auto_element_dlogs.into_iter() {
580 let mut key = M::zeros(auto_d_count * 2, rlwe_n);
581
582 // sample a
583 key.iter_rows_mut().take(auto_d_count).for_each(|ri| {
584 RandomFillUniformInModulus::random_fill(&mut main_prng, &rlwe_q, ri.as_mut())
585 });
586
587 let key_part_b = value.auto_keys.get(&i).unwrap();
588 assert!(key_part_b.dimension() == (auto_d_count, rlwe_n));
589 izip!(
590 key.iter_rows_mut().skip(auto_d_count),
591 key_part_b.iter_rows()
592 )
593 .for_each(|(to_ri, from_ri)| {
594 to_ri.as_mut().copy_from_slice(from_ri.as_ref());
595 });
596
597 // send to evaluation domain
598 key.iter_rows_mut()
599 .for_each(|ri| rlwe_nttop.forward(ri.as_mut()));
600
601 auto_keys.insert(i, key);
602 }
603
604 // rgsw cts
605 let (rlrg_d_a, rlrg_d_b) = value.parameters.rlwe_rgsw_decomposition_count();
606 let rgsw_ct_out = rlrg_d_a.0 * 2 + rlrg_d_b.0 * 2;
607 let rgsw_cts = value
608 .rgsw_cts
609 .iter()
610 .map(|ct_i_in| {
611 assert!(ct_i_in.dimension() == (rgsw_ct_out, rlwe_n));
612 let mut eval_ct_i_out = M::zeros(rgsw_ct_out, rlwe_n);
613
614 izip!(eval_ct_i_out.iter_rows_mut(), ct_i_in.iter_rows()).for_each(
615 |(to_ri, from_ri)| {
616 to_ri.as_mut().copy_from_slice(from_ri.as_ref());
617 rlwe_nttop.forward(to_ri.as_mut());
618 },
619 );
620
621 eval_ct_i_out
622 })
623 .collect_vec();
624
625 // lwe ksk
626 let d_lwe = value.parameters.lwe_decomposition_count().0;
627 let mut lwe_ksk = M::zeros(rlwe_n * d_lwe, lwe_n + 1);
628 izip!(lwe_ksk.iter_rows_mut(), value.lwe_ksk.as_ref().iter()).for_each(
629 |(lwe_i, bi)| {
630 RandomFillUniformInModulus::random_fill(
631 &mut main_prng,
632 &lwe_q,
633 &mut lwe_i.as_mut()[1..],
634 );
635 lwe_i.as_mut()[0] = *bi;
636 },
637 );
638
639 ServerKeyEvaluationDomain {
640 rgsw_cts,
641 galois_keys: auto_keys,
642 lwe_ksk,
643 _phanton: PhantomData,
644 }
645 }
646 }
647
648 impl<M: Matrix, R, N> PbsKey for ServerKeyEvaluationDomain<M, R, N> {
649 type M = M;
650 fn galois_key_for_auto(&self, k: usize) -> &Self::M {
651 self.galois_keys.get(&k).unwrap()
652 }
653 fn rgsw_ct_lwe_si(&self, si: usize) -> &Self::M {
654 &self.rgsw_cts[si]
655 }
656
657 fn lwe_ksk(&self) -> &Self::M {
658 &self.lwe_ksk
659 }
660 }
661}