phantom-zone/src/bool.rs
src/bool.rsRaw
1use std::{
2 cell::RefCell,
3 collections::HashMap,
4 fmt::{Debug, Display},
5 hash::Hash,
6 marker::PhantomData,
7 thread::panicking,
8};
9
10use itertools::{izip, partition, Itertools};
11use num_traits::{FromPrimitive, Num, One, PrimInt, ToPrimitive, WrappingSub, Zero};
12
13use crate::{
14 backend::{ArithmeticOps, ModInit, ModularOpsU64, VectorOps},
15 decomposer::{gadget_vector, Decomposer, DefaultDecomposer, NumInfo},
16 lwe::{decrypt_lwe, encrypt_lwe, lwe_key_switch, lwe_ksk_keygen, measure_noise_lwe, LweSecret},
17 ntt::{Ntt, NttBackendU64, NttInit},
18 random::{DefaultSecureRng, RandomGaussianDist, RandomUniformDist},
19 rgsw::{
20 decrypt_rlwe, encrypt_rgsw, galois_auto, galois_key_gen, generate_auto_map, rlwe_by_rgsw,
21 IsTrivial, RlweCiphertext, RlweSecret,
22 },
23 utils::{generate_prime, mod_exponent, TryConvertFrom, WithLocal},
24 Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut, Secret,
25};
26
27thread_local! {
28 pub(crate) static CLIENT_KEY: RefCell<ClientKey> = RefCell::new(ClientKey::random());
29}
30
31trait PbsKey {
32 type M: Matrix;
33
34 /// RGSW ciphertext of LWE secret elements
35 fn rgsw_ct_lwe_si(&self, si: usize) -> &Self::M;
36 /// Key for automorphism
37 fn galois_key_for_auto(&self, k: isize) -> &Self::M;
38 /// LWE ksk to key switch from RLWE secret to LWE secret
39 fn lwe_ksk(&self) -> &Self::M;
40}
41trait PbsParameters {
42 type Element;
43 type D: Decomposer<Element = Self::Element>;
44 fn rlwe_q(&self) -> Self::Element;
45 fn lwe_q(&self) -> Self::Element;
46 fn br_q(&self) -> usize;
47 fn d_rgsw(&self) -> usize;
48 fn d_lwe(&self) -> usize;
49 fn rlwe_n(&self) -> usize;
50 fn lwe_n(&self) -> usize;
51 /// Embedding fator for ring X^{q}+1 inside
52 fn embedding_factor(&self) -> usize;
53 /// generator g
54 fn g(&self) -> isize;
55 fn decomoposer_lwe(&self) -> &Self::D;
56 fn decomoposer_rlwe(&self) -> &Self::D;
57
58 /// Maps a \in Z^*_{q} to discrete log k, with generator g (i.e. g^k =
59 /// a). Returned vector is of size q that stores dlog of a at `vec[a]`.
60 /// For any a, if k is s.t. a = g^{k}, then k is expressed as k. If k is s.t
61 /// a = -g^{k}, then k is expressed as k=k+q/2
62 fn g_k_dlog_map(&self) -> &[usize];
63 fn rlwe_auto_map(&self, k: isize) -> &(Vec<usize>, Vec<bool>);
64}
65
66#[derive(Clone)]
67struct ClientKey {
68 sk_rlwe: RlweSecret,
69 sk_lwe: LweSecret,
70}
71
72impl ClientKey {
73 fn random() -> Self {
74 let sk_rlwe = RlweSecret::random(0, 0);
75 let sk_lwe = LweSecret::random(0, 0);
76 Self { sk_rlwe, sk_lwe }
77 }
78}
79
80impl WithLocal for ClientKey {
81 fn with_local<F, R>(func: F) -> R
82 where
83 F: Fn(&Self) -> R,
84 {
85 CLIENT_KEY.with_borrow(|client_key| func(client_key))
86 }
87
88 fn with_local_mut<F, R>(func: F) -> R
89 where
90 F: Fn(&mut Self) -> R,
91 {
92 CLIENT_KEY.with_borrow_mut(|client_key| func(client_key))
93 }
94}
95
96fn set_client_key(key: &ClientKey) {
97 ClientKey::with_local_mut(|k| *k = key.clone())
98}
99
100struct ServerKey<M> {
101 /// Rgsw cts of LWE secret elements
102 rgsw_cts: Vec<M>,
103 /// Galois keys
104 galois_keys: HashMap<isize, M>,
105 /// LWE ksk to key switching LWE ciphertext from RLWE secret to LWE secret
106 lwe_ksk: M,
107}
108
109//FIXME(Jay): Figure out a way for BoolEvaluator to have access to ServerKey
110// via a pointer and implement PbsKey for BoolEvaluator instead of ServerKey
111// directly
112impl<M: Matrix> PbsKey for ServerKey<M> {
113 type M = M;
114 fn galois_key_for_auto(&self, k: isize) -> &Self::M {
115 self.galois_keys.get(&k).unwrap()
116 }
117 fn rgsw_ct_lwe_si(&self, si: usize) -> &Self::M {
118 &self.rgsw_cts[si]
119 }
120
121 fn lwe_ksk(&self) -> &Self::M {
122 &self.lwe_ksk
123 }
124}
125
126struct BoolParameters<El> {
127 rlwe_q: El,
128 rlwe_logq: usize,
129 lwe_q: El,
130 lwe_logq: usize,
131 br_q: usize,
132 rlwe_n: usize,
133 lwe_n: usize,
134 d_rgsw: usize,
135 logb_rgsw: usize,
136 d_lwe: usize,
137 logb_lwe: usize,
138 g: usize,
139 w: usize,
140}
141
142struct BoolEvaluator<M, E, Ntt, ModOp>
143where
144 M: Matrix,
145{
146 parameters: BoolParameters<E>,
147 decomposer_rlwe: DefaultDecomposer<E>,
148 decomposer_lwe: DefaultDecomposer<E>,
149 g_k_dlog_map: Vec<usize>,
150 rlwe_nttop: Ntt,
151 rlwe_modop: ModOp,
152 lwe_modop: ModOp,
153 embedding_factor: usize,
154 nand_test_vec: M::R,
155 rlweq_by8: M::MatElement,
156 rlwe_auto_maps: Vec<(Vec<usize>, Vec<bool>)>,
157 scratch_lwen_plus1: M::R,
158 scratch_dplus2_ring: M,
159 _phantom: PhantomData<M>,
160}
161
162impl<M, NttOp, ModOp> BoolEvaluator<M, M::MatElement, NttOp, ModOp>
163where
164 NttOp: NttInit<Element = M::MatElement> + Ntt<Element = M::MatElement>,
165 ModOp: ModInit<Element = M::MatElement>
166 + ArithmeticOps<Element = M::MatElement>
167 + VectorOps<Element = M::MatElement>,
168 M::MatElement: PrimInt + Debug + Display + NumInfo + FromPrimitive + WrappingSub,
169 M: MatrixEntity + MatrixMut,
170 M::R: TryConvertFrom<[i32], Parameters = M::MatElement> + RowEntity,
171 M: TryConvertFrom<[i32], Parameters = M::MatElement>,
172 <M as Matrix>::R: RowMut,
173 DefaultSecureRng: RandomGaussianDist<[M::MatElement], Parameters = M::MatElement>
174 + RandomGaussianDist<M::MatElement, Parameters = M::MatElement>
175 + RandomUniformDist<[M::MatElement], Parameters = M::MatElement>,
176{
177 fn new(parameters: BoolParameters<M::MatElement>) -> Self {
178 //TODO(Jay): Run sanity checks for modulus values in parameters
179 assert!(parameters.br_q.is_power_of_two());
180
181 let decomposer_rlwe =
182 DefaultDecomposer::new(parameters.rlwe_q, parameters.logb_rgsw, parameters.d_rgsw);
183 let decomposer_lwe =
184 DefaultDecomposer::new(parameters.lwe_q, parameters.logb_lwe, parameters.d_lwe);
185
186 // generatr dlog map s.t. g^{k} % q = a, for all a \in Z*_{q}
187 let g = parameters.g;
188 let q = parameters.br_q;
189 let mut g_k_dlog_map = vec![0usize; q];
190 for i in 0..q / 2 {
191 let v = mod_exponent(g as u64, i as u64, q as u64) as usize;
192 // g^i
193 g_k_dlog_map[v] = i;
194 // -(g^i)
195 g_k_dlog_map[q - v] = i + (q / 2);
196 }
197
198 let embedding_factor = (2 * parameters.rlwe_n) / q;
199
200 let rlwe_nttop = NttOp::new(parameters.rlwe_q, parameters.rlwe_n);
201 let rlwe_modop = ModInit::new(parameters.rlwe_q);
202 let lwe_modop = ModInit::new(parameters.lwe_q);
203
204 // set test vectors
205 let el_one = M::MatElement::one();
206 let nand_map = |index: usize, qby8: usize| {
207 if index < (3 * qby8) {
208 true
209 } else {
210 false
211 }
212 };
213
214 let q = parameters.br_q;
215 let qby2 = q >> 1;
216 let qby8 = q >> 3;
217 let qby16 = q >> 4;
218 let mut nand_test_vec = M::R::zeros(qby2);
219 // Q/8 (Q: rlwe_q)
220 let rlwe_qby8 =
221 M::MatElement::from_f64((parameters.rlwe_q.to_f64().unwrap() / 8.0).round()).unwrap();
222 let true_m_el = rlwe_qby8;
223 // -Q/8
224 let false_m_el = parameters.rlwe_q - rlwe_qby8;
225 for i in 0..qby2 {
226 let v = nand_map(i, qby8);
227 if v {
228 nand_test_vec.as_mut()[i] = true_m_el;
229 } else {
230 nand_test_vec.as_mut()[i] = false_m_el;
231 }
232 }
233 // Rotate and negate by q/16
234 let mut tmp = M::R::zeros(qby2);
235 tmp.as_mut()[..qby2 - qby16].copy_from_slice(&nand_test_vec.as_ref()[qby16..]);
236 tmp.as_mut()[qby2 - qby16..].copy_from_slice(&nand_test_vec.as_ref()[..qby16]);
237 tmp.as_mut()[qby2 - qby16..].iter_mut().for_each(|v| {
238 *v = parameters.rlwe_q - *v;
239 });
240 let nand_test_vec = tmp;
241
242 // v(X) -> v(X^{-g})
243 let (auto_map_index, auto_map_sign) = generate_auto_map(qby2, -(g as isize));
244 let mut nand_test_vec_autog = M::R::zeros(qby2);
245 izip!(
246 nand_test_vec.as_ref().iter(),
247 auto_map_index.iter(),
248 auto_map_sign.iter()
249 )
250 .for_each(|(v, to_index, to_sign)| {
251 if !to_sign {
252 // negate
253 nand_test_vec_autog.as_mut()[*to_index] = parameters.rlwe_q - *v;
254 } else {
255 nand_test_vec_autog.as_mut()[*to_index] = *v;
256 }
257 });
258
259 // auto map indices and sign
260 let mut rlwe_auto_maps = vec![];
261 let ring_size = parameters.rlwe_n;
262 let g = parameters.g as isize;
263 for i in [g, -g] {
264 rlwe_auto_maps.push(generate_auto_map(ring_size, i))
265 }
266
267 // create srcatch spaces
268 let scratch_lwen_plus1 = M::R::zeros(parameters.lwe_n + 1);
269 let scratch_dplus2_ring = M::zeros(parameters.d_rgsw + 2, parameters.rlwe_n);
270
271 BoolEvaluator {
272 parameters: parameters,
273 decomposer_lwe,
274 decomposer_rlwe,
275 g_k_dlog_map,
276 embedding_factor,
277 lwe_modop,
278 rlwe_modop,
279 rlwe_nttop,
280 nand_test_vec: nand_test_vec_autog,
281 rlweq_by8: rlwe_qby8,
282 rlwe_auto_maps,
283 scratch_lwen_plus1,
284 scratch_dplus2_ring,
285 _phantom: PhantomData,
286 }
287 }
288
289 fn client_key(&self) -> ClientKey {
290 let sk_lwe = LweSecret::random(self.parameters.lwe_n >> 1, self.parameters.lwe_n);
291 let sk_rlwe = RlweSecret::random(self.parameters.rlwe_n >> 1, self.parameters.rlwe_n);
292 ClientKey { sk_rlwe, sk_lwe }
293 }
294
295 fn server_key(&self, client_key: &ClientKey) -> ServerKey<M> {
296 let sk_rlwe = &client_key.sk_rlwe;
297 let sk_lwe = &client_key.sk_lwe;
298
299 let d_rgsw_gadget_vec = gadget_vector(
300 self.parameters.rlwe_logq,
301 self.parameters.logb_rgsw,
302 self.parameters.d_rgsw,
303 );
304
305 // generate galois key -g, g
306 let mut galois_keys = HashMap::new();
307 let g = self.parameters.g as isize;
308 for i in [g, -g] {
309 let gk = DefaultSecureRng::with_local_mut(|rng| {
310 let mut ksk_out = M::zeros(self.parameters.d_rgsw * 2, self.parameters.rlwe_n);
311 galois_key_gen(
312 &mut ksk_out,
313 sk_rlwe,
314 i,
315 &d_rgsw_gadget_vec,
316 &self.rlwe_modop,
317 &self.rlwe_nttop,
318 rng,
319 );
320 ksk_out
321 });
322
323 galois_keys.insert(i, gk);
324 }
325
326 // generate rgsw ciphertexts RGSW(si) where si is i^th LWE secret element
327 let ring_size = self.parameters.rlwe_n;
328 let rlwe_q = self.parameters.rlwe_q;
329 let rgsw_cts = sk_lwe
330 .values()
331 .iter()
332 .map(|si| {
333 // X^{si}; assume |emebedding_factor * si| < N
334 let mut m = M::zeros(1, ring_size);
335 let si = (self.embedding_factor as i32) * si;
336 // dbg!(si);
337 if si < 0 {
338 // X^{-i} = X^{2N - i} = -X^{N-i}
339 m.set(
340 0,
341 ring_size - (si.abs() as usize),
342 rlwe_q - M::MatElement::one(),
343 );
344 } else {
345 // X^{i}
346 m.set(0, (si.abs() as usize), M::MatElement::one());
347 }
348 self.rlwe_nttop.forward(m.get_row_mut(0));
349
350 let rgsw_si = DefaultSecureRng::with_local_mut(|rng| {
351 let mut rgsw_si = M::zeros(self.parameters.d_rgsw * 4, ring_size);
352 encrypt_rgsw(
353 &mut rgsw_si,
354 &m,
355 &d_rgsw_gadget_vec,
356 sk_rlwe,
357 &self.rlwe_modop,
358 &self.rlwe_nttop,
359 rng,
360 );
361 rgsw_si
362 });
363 rgsw_si
364 })
365 .collect_vec();
366
367 // LWE KSK from RLWE secret s -> LWE secret z
368 let d_lwe_gadget = gadget_vector(
369 self.parameters.lwe_logq,
370 self.parameters.logb_lwe,
371 self.parameters.d_lwe,
372 );
373 let mut lwe_ksk = DefaultSecureRng::with_local_mut(|rng| {
374 let mut out = M::zeros(self.parameters.d_lwe * ring_size, self.parameters.lwe_n + 1);
375 lwe_ksk_keygen(
376 &sk_rlwe.values(),
377 &sk_lwe.values(),
378 &mut out,
379 &d_lwe_gadget,
380 &self.lwe_modop,
381 rng,
382 );
383 out
384 });
385
386 ServerKey {
387 rgsw_cts,
388 galois_keys,
389 lwe_ksk,
390 }
391 }
392
393 /// TODO(Jay): Fetch client key from thread local
394 pub fn encrypt(&self, m: bool, client_key: &ClientKey) -> M::R {
395 let m = if m {
396 // Q/8
397 self.rlweq_by8
398 } else {
399 // -Q/8
400 self.parameters.rlwe_q - self.rlweq_by8
401 };
402
403 DefaultSecureRng::with_local_mut(|rng| {
404 let mut lwe_out = M::R::zeros(self.parameters.rlwe_n + 1);
405 encrypt_lwe(
406 &mut lwe_out,
407 &m,
408 client_key.sk_rlwe.values(),
409 &self.rlwe_modop,
410 rng,
411 );
412 lwe_out
413 })
414 }
415
416 pub fn decrypt(&self, lwe_ct: &M::R, client_key: &ClientKey) -> bool {
417 let m = decrypt_lwe(lwe_ct, client_key.sk_rlwe.values(), &self.rlwe_modop);
418 let m = {
419 // m + q/8 => {0,q/4 1}
420 (((m + self.rlweq_by8).to_f64().unwrap() * 4.0)
421 / self.parameters.rlwe_q.to_f64().unwrap())
422 .round()
423 .to_usize()
424 .unwrap()
425 % 4
426 };
427
428 if m == 0 {
429 false
430 } else if m == 1 {
431 true
432 } else {
433 panic!("Incorrect bool decryption. Got m={m} but expected m to be 0 or 1")
434 }
435 }
436
437 pub fn nand(
438 &self,
439 c0: &M::R,
440 c1: &M::R,
441 server_key: &ServerKey<M>,
442 scratch_lwen_plus1: &mut M::R,
443 scratch_matrix_dplus2_ring: &mut M,
444 ) -> M::R {
445 // ClientKey::with_local(|ck| {
446 // let c0_noise = measure_noise_lwe(
447 // c0,
448 // ck.sk_rlwe.values(),
449 // &self.rlwe_modop,
450 // &(self.rlwe_q() - self.rlweq_by8),
451 // );
452 // let c1_noise =
453 // measure_noise_lwe(c1, ck.sk_rlwe.values(), &self.rlwe_modop,
454 // &(self.rlweq_by8)); println!("c0 noise: {c0_noise}; c1 noise:
455 // {c1_noise}"); });
456
457 let mut c_out = M::R::zeros(c0.as_ref().len());
458 let modop = &self.rlwe_modop;
459 izip!(
460 c_out.as_mut().iter_mut(),
461 c0.as_ref().iter(),
462 c1.as_ref().iter()
463 )
464 .for_each(|(o, i0, i1)| {
465 *o = modop.add(i0, i1);
466 });
467 // +Q/8
468 c_out.as_mut()[0] = modop.add(&c_out.as_ref()[0], &self.rlweq_by8);
469
470 // ClientKey::with_local(|ck| {
471 // let noise = measure_noise_lwe(
472 // &c_out,
473 // ck.sk_rlwe.values(),
474 // &self.rlwe_modop,
475 // &(self.rlweq_by8),
476 // );
477 // println!("cout_noise: {noise}");
478 // });
479
480 // PBS
481 pbs(
482 self,
483 &self.nand_test_vec,
484 &mut c_out,
485 scratch_lwen_plus1,
486 scratch_matrix_dplus2_ring,
487 &self.lwe_modop,
488 &self.rlwe_modop,
489 &self.rlwe_nttop,
490 server_key,
491 );
492
493 c_out
494 }
495}
496
497impl<M: Matrix, NttOp, ModOp> PbsParameters for BoolEvaluator<M, M::MatElement, NttOp, ModOp>
498where
499 M::MatElement: PrimInt + WrappingSub + Debug,
500{
501 type Element = M::MatElement;
502 type D = DefaultDecomposer<M::MatElement>;
503 fn rlwe_auto_map(&self, k: isize) -> &(Vec<usize>, Vec<bool>) {
504 let g = self.parameters.g as isize;
505 if k == g {
506 &self.rlwe_auto_maps[0]
507 } else if k == -g {
508 &self.rlwe_auto_maps[1]
509 } else {
510 panic!("RLWE auto map only supports k in [-g, g], but got k={k}");
511 }
512 }
513
514 fn br_q(&self) -> usize {
515 self.parameters.br_q
516 }
517 fn d_lwe(&self) -> usize {
518 self.parameters.d_lwe
519 }
520 fn d_rgsw(&self) -> usize {
521 self.parameters.d_rgsw
522 }
523 fn decomoposer_lwe(&self) -> &Self::D {
524 &self.decomposer_lwe
525 }
526 fn decomoposer_rlwe(&self) -> &Self::D {
527 &self.decomposer_rlwe
528 }
529 fn embedding_factor(&self) -> usize {
530 self.embedding_factor
531 }
532 fn g(&self) -> isize {
533 self.parameters.g as isize
534 }
535 fn g_k_dlog_map(&self) -> &[usize] {
536 &self.g_k_dlog_map
537 }
538 fn lwe_n(&self) -> usize {
539 self.parameters.lwe_n
540 }
541 fn lwe_q(&self) -> Self::Element {
542 self.parameters.lwe_q
543 }
544 fn rlwe_n(&self) -> usize {
545 self.parameters.rlwe_n
546 }
547 fn rlwe_q(&self) -> Self::Element {
548 self.parameters.rlwe_q
549 }
550}
551
552/// LMKCY+ Blind rotation
553///
554/// gk_to_si: [-g^0, -g^1, .., -g^{q/2-1}, g^0, ..., g^{q/2-1}]
555fn blind_rotation<
556 MT: IsTrivial + MatrixMut,
557 Mmut: MatrixMut<MatElement = MT::MatElement> + Matrix,
558 D: Decomposer<Element = MT::MatElement>,
559 NttOp: Ntt<Element = MT::MatElement>,
560 ModOp: ArithmeticOps<Element = MT::MatElement> + VectorOps<Element = MT::MatElement>,
561 K: PbsKey<M = Mmut>,
562 P: PbsParameters<Element = MT::MatElement>,
563>(
564 trivial_rlwe_test_poly: &mut MT,
565 scratch_matrix_dplus2_ring: &mut Mmut,
566 g: isize,
567 w: usize,
568 q: usize,
569 gk_to_si: &[Vec<usize>],
570 decomposer: &D,
571 ntt_op: &NttOp,
572 mod_op: &ModOp,
573 parameters: &P,
574 pbs_key: &K,
575) where
576 <Mmut as Matrix>::R: RowMut,
577 Mmut::MatElement: Copy + Zero,
578 <MT as Matrix>::R: RowMut,
579{
580 let q_by_2 = q / 2;
581
582 // -(g^k)
583 for i in (1..q_by_2).rev() {
584 gk_to_si[q_by_2 + i].iter().for_each(|s_index| {
585 rlwe_by_rgsw(
586 trivial_rlwe_test_poly,
587 pbs_key.rgsw_ct_lwe_si(*s_index),
588 scratch_matrix_dplus2_ring,
589 decomposer,
590 ntt_op,
591 mod_op,
592 );
593 });
594
595 let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(g);
596 galois_auto(
597 trivial_rlwe_test_poly,
598 pbs_key.galois_key_for_auto(g),
599 scratch_matrix_dplus2_ring,
600 &auto_map_index,
601 &auto_map_sign,
602 mod_op,
603 ntt_op,
604 decomposer,
605 );
606 }
607
608 // -(g^0)
609 gk_to_si[q_by_2].iter().for_each(|s_index| {
610 rlwe_by_rgsw(
611 trivial_rlwe_test_poly,
612 pbs_key.rgsw_ct_lwe_si(*s_index),
613 scratch_matrix_dplus2_ring,
614 decomposer,
615 ntt_op,
616 mod_op,
617 );
618 });
619 let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(-g);
620 galois_auto(
621 trivial_rlwe_test_poly,
622 pbs_key.galois_key_for_auto(-g),
623 scratch_matrix_dplus2_ring,
624 &auto_map_index,
625 &auto_map_sign,
626 mod_op,
627 ntt_op,
628 decomposer,
629 );
630
631 // +(g^k)
632 for i in (1..q_by_2).rev() {
633 gk_to_si[i].iter().for_each(|s_index| {
634 rlwe_by_rgsw(
635 trivial_rlwe_test_poly,
636 pbs_key.rgsw_ct_lwe_si(*s_index),
637 scratch_matrix_dplus2_ring,
638 decomposer,
639 ntt_op,
640 mod_op,
641 );
642 });
643
644 let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(g);
645 galois_auto(
646 trivial_rlwe_test_poly,
647 pbs_key.galois_key_for_auto(g),
648 scratch_matrix_dplus2_ring,
649 &auto_map_index,
650 &auto_map_sign,
651 mod_op,
652 ntt_op,
653 decomposer,
654 );
655 }
656
657 // +(g^0)
658 gk_to_si[0].iter().for_each(|s_index| {
659 rlwe_by_rgsw(
660 trivial_rlwe_test_poly,
661 pbs_key.rgsw_ct_lwe_si(gk_to_si[q_by_2][*s_index]),
662 scratch_matrix_dplus2_ring,
663 decomposer,
664 ntt_op,
665 mod_op,
666 );
667 });
668}
669
670/// - Mod down
671/// - key switching
672/// - mod down
673/// - blind rotate
674fn pbs<
675 M: Matrix + MatrixMut + MatrixEntity,
676 P: PbsParameters<Element = M::MatElement>,
677 NttOp: Ntt<Element = M::MatElement>,
678 ModOp: ArithmeticOps<Element = M::MatElement> + VectorOps<Element = M::MatElement>,
679 K: PbsKey<M = M>,
680>(
681 parameters: &P,
682 test_vec: &M::R,
683 lwe_in: &mut M::R,
684 scratch_lwen_plus1: &mut M::R,
685 scratch_matrix_dplus2_ring: &mut M,
686 modop_lweq: &ModOp,
687 modop_rlweq: &ModOp,
688 nttop_rlweq: &NttOp,
689 pbs_key: &K,
690) where
691 // FIXME(Jay): TryConvertFrom<[i32], Parameters = M::MatElement> are only needed for
692 // debugging purposes
693 <M as Matrix>::R: RowMut + TryConvertFrom<[i32], Parameters = M::MatElement>,
694 M::MatElement: PrimInt + ToPrimitive + FromPrimitive + One + Copy + Zero + Display,
695{
696 let rlwe_q = parameters.rlwe_q();
697 let lwe_q = parameters.lwe_q();
698 let br_q = parameters.br_q();
699 let rlwe_qf64 = rlwe_q.to_f64().unwrap();
700 let lwe_qf64 = lwe_q.to_f64().unwrap();
701 let br_qf64 = br_q.to_f64().unwrap();
702 let rlwe_n = parameters.rlwe_n();
703
704 // moddown Q -> Q_ks
705 lwe_in.as_mut().iter_mut().for_each(|v| {
706 *v =
707 M::MatElement::from_f64(((v.to_f64().unwrap() * lwe_qf64) / rlwe_qf64).round()).unwrap()
708 });
709
710 PBSTracer::with_local_mut(|t| {
711 let out = lwe_in
712 .as_ref()
713 .iter()
714 .map(|v| v.to_u64().unwrap())
715 .collect_vec();
716 t.ct_lwe_q_mod = out;
717 });
718
719 // key switch RLWE secret to LWE secret
720 scratch_lwen_plus1.as_mut().fill(M::MatElement::zero());
721 lwe_key_switch(
722 scratch_lwen_plus1,
723 lwe_in,
724 pbs_key.lwe_ksk(),
725 modop_lweq,
726 parameters.decomoposer_lwe(),
727 );
728
729 PBSTracer::with_local_mut(|t| {
730 let out = scratch_lwen_plus1
731 .as_ref()
732 .iter()
733 .map(|v| v.to_u64().unwrap())
734 .collect_vec();
735 t.ct_lwe_q_mod_after_ksk = out;
736 });
737
738 // odd mowdown Q_ks -> q
739 let g_k_dlog_map = parameters.g_k_dlog_map();
740 let mut g_k_si = vec![vec![]; br_q];
741 scratch_lwen_plus1
742 .as_ref()
743 .iter()
744 .skip(1)
745 .enumerate()
746 .for_each(|(index, v)| {
747 let odd_v = mod_switch_odd(v.to_f64().unwrap(), lwe_qf64, br_qf64);
748 let k = g_k_dlog_map[odd_v];
749 g_k_si[k].push(index);
750 });
751
752 PBSTracer::with_local_mut(|t| {
753 let out = scratch_lwen_plus1
754 .as_ref()
755 .iter()
756 .map(|v| mod_switch_odd(v.to_f64().unwrap(), lwe_qf64, br_qf64) as u64)
757 .collect_vec();
758 t.ct_br_q_mod = out;
759 });
760
761 // handle b and set trivial test RLWE
762 let g = parameters.g() as usize;
763 let g_times_b = (g * mod_switch_odd(
764 scratch_lwen_plus1.as_ref()[0].to_f64().unwrap(),
765 lwe_qf64,
766 br_qf64,
767 )) % (br_q);
768 // v = (v(X) * X^{g*b}) mod X^{q/2}+1
769 let br_qby2 = br_q / 2;
770 let mut gb_monomial_sign = true;
771 let mut gb_monomial_exp = g_times_b;
772 // X^{g*b} mod X^{q/2}+1
773 if gb_monomial_exp > br_qby2 {
774 gb_monomial_exp -= br_qby2;
775 gb_monomial_sign = false
776 }
777 // monomial mul
778 let mut trivial_rlwe_test_poly = RlweCiphertext(M::zeros(2, rlwe_n), true);
779 if parameters.embedding_factor() == 1 {
780 monomial_mul(
781 test_vec.as_ref(),
782 trivial_rlwe_test_poly.get_row_mut(1).as_mut(),
783 gb_monomial_exp,
784 gb_monomial_sign,
785 br_qby2,
786 modop_rlweq,
787 );
788 } else {
789 // use lwe_in to store the `t = v(X) * X^{g*2} mod X^{q/2}+1` temporarily. This
790 // works because q/2 <= N (where N is lwe_in LWE dimension) always.
791 monomial_mul(
792 test_vec.as_ref(),
793 &mut lwe_in.as_mut()[..br_qby2],
794 gb_monomial_exp,
795 gb_monomial_sign,
796 br_qby2,
797 modop_rlweq,
798 );
799
800 // emebed poly `t` in ring X^{q/2}+1 inside the bigger ring X^{N}+1
801 let embed_factor = parameters.embedding_factor();
802 let partb_trivial_rlwe = trivial_rlwe_test_poly.get_row_mut(1);
803 lwe_in.as_ref()[..br_qby2]
804 .iter()
805 .enumerate()
806 .for_each(|(index, v)| {
807 partb_trivial_rlwe[embed_factor * index] = *v;
808 });
809 }
810
811 // blind rotate
812 blind_rotation(
813 &mut trivial_rlwe_test_poly,
814 scratch_matrix_dplus2_ring,
815 parameters.g(),
816 1,
817 br_q,
818 &g_k_si,
819 parameters.decomoposer_rlwe(),
820 nttop_rlweq,
821 modop_rlweq,
822 parameters,
823 pbs_key,
824 );
825
826 // ClientKey::with_local(|ck| {
827 // let ring_size = parameters.rlwe_n();
828 // let mut rlwe_ct = vec![vec![0u64; ring_size]; 2];
829 // izip!(
830 // rlwe_ct[0].iter_mut(),
831 // trivial_rlwe_test_poly.0.get_row_slice(0)
832 // )
833 // .for_each(|(t, f)| {
834 // *t = f.to_u64().unwrap();
835 // });
836 // izip!(
837 // rlwe_ct[1].iter_mut(),
838 // trivial_rlwe_test_poly.0.get_row_slice(1)
839 // )
840 // .for_each(|(t, f)| {
841 // *t = f.to_u64().unwrap();
842 // });
843 // let mut m_out = vec![vec![0u64; ring_size]];
844 // let modop = ModularOpsU64::new(rlwe_q.to_u64().unwrap());
845 // let nttop = NttBackendU64::new(rlwe_q.to_u64().unwrap(), ring_size);
846 // decrypt_rlwe(&rlwe_ct, ck.sk_rlwe.values(), &mut m_out, &nttop, &modop);
847
848 // println!("RLWE post PBS message: {:?}", m_out[0]);
849 // });
850
851 // sample extract
852 sample_extract(lwe_in, &trivial_rlwe_test_poly, modop_rlweq, 0);
853}
854
855fn mod_switch_odd(v: f64, from_q: f64, to_q: f64) -> usize {
856 // println!("v: {v}, odd_v: {odd_v}, lwe_q:{lwe_q}, br_q:{br_q}");
857 let odd_v = (((v * to_q) / (from_q)).floor()).to_usize().unwrap();
858 // println!(
859 // "v: {v}, odd_v: {odd_v}, returned_oddv: {},lwe_q:{from_q}, br_q:{to_q}",
860 // odd_v + ((odd_v & 1) ^ 1)
861 // );
862 //TODO(Jay): check correctness of this
863 odd_v + ((odd_v & 1) ^ 1)
864}
865
866fn sample_extract<M: Matrix + MatrixMut, ModOp: ArithmeticOps<Element = M::MatElement>>(
867 lwe_out: &mut M::R,
868 rlwe_in: &M,
869 mod_op: &ModOp,
870 index: usize,
871) where
872 <M as Matrix>::R: RowMut,
873 M::MatElement: Copy,
874{
875 let ring_size = rlwe_in.dimension().1;
876
877 // index..=0
878 let to = &mut lwe_out.as_mut()[1..];
879 let from = rlwe_in.get_row_slice(0);
880 for i in 0..index + 1 {
881 to[i] = from[index - i];
882 }
883
884 // -(N..index)
885 for i in index + 1..ring_size {
886 to[i] = mod_op.neg(&from[ring_size + index - i]);
887 }
888
889 // set b
890 lwe_out.as_mut()[0] = *rlwe_in.get(1, index);
891}
892
893/// TODO(Jay): Write tests for monomial mul
894fn monomial_mul<El, ModOp: ArithmeticOps<Element = El>>(
895 p_in: &[El],
896 p_out: &mut [El],
897 mon_exp: usize,
898 mon_sign: bool,
899 ring_size: usize,
900 mod_op: &ModOp,
901) where
902 El: Copy,
903{
904 debug_assert!(p_in.as_ref().len() == ring_size);
905 debug_assert!(p_in.as_ref().len() == p_out.as_ref().len());
906 debug_assert!(mon_exp < ring_size);
907
908 p_in.as_ref().iter().enumerate().for_each(|(index, v)| {
909 let mut to_index = index + mon_exp;
910 let mut to_sign = mon_sign;
911 if to_index >= ring_size {
912 to_index = to_index - ring_size;
913 to_sign = !to_sign;
914 }
915
916 if !to_sign {
917 p_out.as_mut()[to_index] = mod_op.neg(v);
918 } else {
919 p_out.as_mut()[to_index] = *v;
920 }
921 });
922}
923
924thread_local! {
925 static PBS_TRACER: RefCell<PBSTracer<Vec<Vec<u64>>>> = RefCell::new(PBSTracer::default());
926}
927
928#[derive(Default)]
929struct PBSTracer<M>
930where
931 M: Matrix + Default,
932{
933 pub(crate) ct_lwe_q_mod: M::R,
934 pub(crate) ct_lwe_q_mod_after_ksk: M::R,
935 pub(crate) ct_br_q_mod: Vec<u64>,
936}
937
938impl PBSTracer<Vec<Vec<u64>>> {
939 fn trace(&self, parameters: &BoolParameters<u64>, client_key: &ClientKey, expected_m: bool) {
940 let lwe_q = parameters.lwe_q;
941 let lwe_qby8 = ((lwe_q as f64) / 8.0).round() as u64;
942 let expected_m_lweq = if expected_m {
943 lwe_qby8
944 } else {
945 lwe_q - lwe_qby8
946 };
947 let modop_lweq = ModularOpsU64::new(lwe_q);
948 // noise after mod down Q -> Q_ks
949 let noise0 = {
950 measure_noise_lwe(
951 &self.ct_lwe_q_mod,
952 client_key.sk_rlwe.values(),
953 &modop_lweq,
954 &expected_m_lweq,
955 )
956 };
957
958 // noise after key switch from RLWE -> LWE
959 let noise1 = {
960 measure_noise_lwe(
961 &self.ct_lwe_q_mod_after_ksk,
962 client_key.sk_lwe.values(),
963 &modop_lweq,
964 &expected_m_lweq,
965 )
966 };
967
968 // noise after mod down odd from Q_ks -> q
969 let br_q = parameters.br_q as u64;
970 let expected_m_brq = if expected_m {
971 br_q >> 3
972 } else {
973 br_q - (br_q >> 3)
974 };
975 let modop_br_q = ModularOpsU64::new(br_q);
976 let noise2 = {
977 measure_noise_lwe(
978 &self.ct_br_q_mod,
979 client_key.sk_lwe.values(),
980 &modop_br_q,
981 &expected_m_brq,
982 )
983 };
984
985 println!(
986 "
987 m: {expected_m},
988 Noise after mod down Q -> Q_ks: {noise0},
989 Noise after key switch from RLWE -> LWE: {noise1},
990 Noise after mod dwon Q_ks -> q: {noise2}
991 "
992 );
993 }
994}
995
996impl WithLocal for PBSTracer<Vec<Vec<u64>>> {
997 fn with_local<F, R>(func: F) -> R
998 where
999 F: Fn(&Self) -> R,
1000 {
1001 PBS_TRACER.with_borrow(|t| func(t))
1002 }
1003
1004 fn with_local_mut<F, R>(func: F) -> R
1005 where
1006 F: Fn(&mut Self) -> R,
1007 {
1008 PBS_TRACER.with_borrow_mut(|t| func(t))
1009 }
1010}
1011
1012#[cfg(test)]
1013mod tests {
1014 use crate::{backend::ModularOpsU64, ntt::NttBackendU64, random::DEFAULT_RNG};
1015
1016 use super::*;
1017
1018 const SP_BOOL_PARAMS: BoolParameters<u64> = BoolParameters::<u64> {
1019 rlwe_q: 268369921u64,
1020 rlwe_logq: 28,
1021 lwe_q: 1 << 16,
1022 lwe_logq: 16,
1023 br_q: 1 << 10,
1024 rlwe_n: 1 << 10,
1025 lwe_n: 493,
1026 d_rgsw: 3,
1027 logb_rgsw: 8,
1028 d_lwe: 3,
1029 logb_lwe: 4,
1030 g: 5,
1031 w: 1,
1032 };
1033
1034 // #[test]
1035 // fn trial() {
1036 // dbg!(generate_prime(28, 1 << 11, 1 << 28));
1037 // }
1038
1039 #[test]
1040 fn encrypt_decrypt_works() {
1041 // let prime = generate_prime(32, 2 * 1024, 1 << 32);
1042 // dbg!(prime);
1043 let bool_evaluator =
1044 BoolEvaluator::<Vec<Vec<u64>>, u64, NttBackendU64, ModularOpsU64>::new(SP_BOOL_PARAMS);
1045 let client_key = bool_evaluator.client_key();
1046 // let sever_key = bool_evaluator.server_key(&client_key);
1047
1048 let mut m = true;
1049 for _ in 0..1000 {
1050 let lwe_ct = bool_evaluator.encrypt(m, &client_key);
1051 let m_back = bool_evaluator.decrypt(&lwe_ct, &client_key);
1052 assert_eq!(m, m_back);
1053 m = !m;
1054 }
1055 }
1056
1057 #[test]
1058 fn trial12() {
1059 // DefaultSecureRng::with_local_mut(|r| {
1060 // let rng = DefaultSecureRng::new_seeded([19u8; 32]);
1061 // *r = rng;
1062 // });
1063
1064 let bool_evaluator =
1065 BoolEvaluator::<Vec<Vec<u64>>, u64, NttBackendU64, ModularOpsU64>::new(SP_BOOL_PARAMS);
1066 // println!("{:?}", bool_evaluator.nand_test_vec);
1067 let client_key = bool_evaluator.client_key();
1068 set_client_key(&client_key);
1069
1070 let server_key = bool_evaluator.server_key(&client_key);
1071
1072 let mut scratch_lwen_plus1 = vec![0u64; bool_evaluator.parameters.lwe_n + 1];
1073 let mut scratch_matrix_dplus2_ring = vec![
1074 vec![0u64; bool_evaluator.parameters.rlwe_n];
1075 bool_evaluator.parameters.d_rgsw + 2
1076 ];
1077
1078 let mut m0 = false;
1079 let mut m1 = true;
1080 let mut ct0 = bool_evaluator.encrypt(m0, &client_key);
1081 let mut ct1 = bool_evaluator.encrypt(m1, &client_key);
1082 for _ in 0..4 {
1083 let ct_back = bool_evaluator.nand(
1084 &ct0,
1085 &ct1,
1086 &server_key,
1087 &mut scratch_lwen_plus1,
1088 &mut scratch_matrix_dplus2_ring,
1089 );
1090
1091 let m_out = !(m0 && m1);
1092
1093 // Trace and measure PBS noise
1094 {
1095 // Trace PBS
1096 PBSTracer::with_local(|t| t.trace(&SP_BOOL_PARAMS, &client_key, m_out));
1097
1098 // Calculate nosie in ciphertext post PBS
1099 let ideal = if m_out {
1100 bool_evaluator.rlweq_by8
1101 } else {
1102 bool_evaluator.rlwe_q() - bool_evaluator.rlweq_by8
1103 };
1104 let noise = measure_noise_lwe(
1105 &ct_back,
1106 client_key.sk_rlwe.values(),
1107 &bool_evaluator.rlwe_modop,
1108 &ideal,
1109 );
1110 println!("PBS noise: {noise}");
1111 }
1112 let m_back = bool_evaluator.decrypt(&ct_back, &client_key);
1113 assert_eq!(m_out, m_back);
1114 println!("----------");
1115
1116 m1 = m0;
1117 m0 = m_out;
1118
1119 ct1 = ct0;
1120 ct0 = ct_back;
1121 }
1122 }
1123}