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