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