phantom-zone/src/bool.rs
src/bool.rsRaw
1use std::{collections::HashMap, fmt::Debug, marker::PhantomData};
2
3use itertools::Itertools;
4use num_traits::{FromPrimitive, Num, One, PrimInt, ToPrimitive, Zero};
5
6use crate::{
7 backend::{ArithmeticOps, ModInit, VectorOps},
8 decomposer::{gadget_vector, Decomposer, DefaultDecomposer, NumInfo},
9 lwe::{decrypt_lwe, encrypt_lwe, lwe_key_switch, lwe_ksk_keygen, LweSecret},
10 ntt::{Ntt, NttInit},
11 random::{DefaultSecureRng, RandomGaussianDist, RandomUniformDist},
12 rgsw::{encrypt_rgsw, galois_auto, galois_key_gen, rlwe_by_rgsw, IsTrivial, RlweSecret},
13 utils::{generate_prime, mod_exponent, TryConvertFrom, WithLocal},
14 Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut, Secret,
15};
16
17trait PbsKey {
18 type M: Matrix;
19
20 fn rgsw_ct_secret_el(&self, si: usize) -> &Self::M;
21 fn galois_key_for_auto(&self, k: isize) -> &Self::M;
22 fn auto_map_index(&self, k: isize) -> &[usize];
23 fn auto_map_sign(&self, k: isize) -> &[bool];
24}
25trait Parameters {
26 type Element;
27 type D: Decomposer<Element = Self::Element>;
28 fn rlwe_q(&self) -> Self::Element;
29 fn lwe_q(&self) -> Self::Element;
30 fn br_q(&self) -> usize;
31 fn d_rgsw(&self) -> usize;
32 fn d_lwe(&self) -> usize;
33 fn rlwe_n(&self) -> usize;
34 fn lwe_n(&self) -> usize;
35 /// Embedding fator for ring X^{q}+1 inside
36 fn embedding_factor(&self) -> usize;
37 /// generator g
38 fn g(&self) -> isize;
39 fn decomoposer_lwe(&self) -> &Self::D;
40 fn decomoposer_rlwe(&self) -> &Self::D;
41
42 /// Maps a \in Z^*_{q} to discrete log k, with generator g (i.e. g^k =
43 /// a). Returned vector is of size q that stores dlog of a at `vec[a]`.
44 /// For any a, if k is s.t. a = g^{k}, then k is expressed as k. If k is s.t
45 /// a = -g^{k}, then k is expressed as k=k+q/2
46 fn g_k_dlog_map(&self) -> &[usize];
47}
48struct ClientKey {
49 sk_rlwe: RlweSecret,
50 sk_lwe: LweSecret,
51}
52
53struct ServerKey<M> {
54 /// Rgsw cts of LWE secret elements
55 rgsw_cts: Vec<M>,
56 /// Galois keys
57 galois_keys: HashMap<isize, M>,
58 /// LWE ksk to key switching LWE ciphertext from RLWE secret to LWE secret
59 lwe_ksk: M,
60}
61
62struct BoolParameters<El> {
63 rlwe_q: El,
64 rlwe_logq: usize,
65 lwe_q: El,
66 lwe_logq: usize,
67 br_q: usize,
68 rlwe_n: usize,
69 lwe_n: usize,
70 d_rgsw: usize,
71 logb_rgsw: usize,
72 d_lwe: usize,
73 logb_lwe: usize,
74 g: usize,
75 w: usize,
76}
77
78struct BoolEvaluator<M, E, Ntt, ModOp> {
79 parameters: BoolParameters<E>,
80 decomposer_rlwe: DefaultDecomposer<E>,
81 decomposer_lwe: DefaultDecomposer<E>,
82 g_k_dlog_map: Vec<usize>,
83 rlwe_nttop: Ntt,
84 rlwe_modop: ModOp,
85 lwe_modop: ModOp,
86 embedding_factor: usize,
87
88 _phantom: PhantomData<M>,
89}
90
91impl<M, NttOp, ModOp> BoolEvaluator<M, M::MatElement, NttOp, ModOp>
92where
93 NttOp: NttInit<Element = M::MatElement> + Ntt<Element = M::MatElement>,
94 ModOp: ModInit<Element = M::MatElement>
95 + ArithmeticOps<Element = M::MatElement>
96 + VectorOps<Element = M::MatElement>,
97 M::MatElement: PrimInt + Debug + NumInfo + FromPrimitive,
98 M: MatrixEntity + MatrixMut,
99 M::R: TryConvertFrom<[i32], Parameters = M::MatElement> + RowEntity,
100 M: TryConvertFrom<[i32], Parameters = M::MatElement>,
101 <M as Matrix>::R: RowMut,
102 DefaultSecureRng: RandomGaussianDist<[M::MatElement], Parameters = M::MatElement>
103 + RandomGaussianDist<M::MatElement, Parameters = M::MatElement>
104 + RandomUniformDist<[M::MatElement], Parameters = M::MatElement>,
105{
106 fn new(parameters: BoolParameters<M::MatElement>) -> Self {
107 //TODO(Jay): Run sanity checks for modulus values in parameters
108
109 let decomposer_rlwe =
110 DefaultDecomposer::new(parameters.rlwe_q, parameters.logb_rgsw, parameters.d_rgsw);
111 let decomposer_lwe =
112 DefaultDecomposer::new(parameters.lwe_q, parameters.logb_lwe, parameters.d_lwe);
113
114 // generatr dlog map s.t. g^{k} % q = a, for all a \in Z*_{q}
115 let g = parameters.g;
116 let q = parameters.br_q;
117 let mut g_k_dlog_map = vec![0usize; q];
118 for i in 0..q / 2 {
119 let v = mod_exponent(g as u64, i as u64, q as u64) as usize;
120 // g^i
121 g_k_dlog_map[v] = i;
122 // -(g^i)
123 g_k_dlog_map[q - v] = i + (q / 2);
124 }
125
126 let embedding_factor = (2 * parameters.rlwe_n) / q;
127
128 let rlwe_nttop = NttOp::new(parameters.rlwe_q, parameters.rlwe_n);
129 let rlwe_modop = ModInit::new(parameters.rlwe_q);
130 let lwe_modop = ModInit::new(parameters.lwe_q);
131
132 BoolEvaluator {
133 parameters: parameters,
134 decomposer_lwe,
135 decomposer_rlwe,
136 g_k_dlog_map,
137 embedding_factor,
138 lwe_modop,
139 rlwe_modop,
140 rlwe_nttop,
141
142 _phantom: PhantomData,
143 }
144 }
145
146 fn client_key(&self) -> ClientKey {
147 let sk_lwe = LweSecret::random(self.parameters.lwe_n >> 1, self.parameters.lwe_n);
148 let sk_rlwe = RlweSecret::random(self.parameters.rlwe_n >> 1, self.parameters.rlwe_n);
149 ClientKey { sk_rlwe, sk_lwe }
150 }
151
152 fn server_key(&self, client_key: &ClientKey) -> ServerKey<M> {
153 let sk_rlwe = &client_key.sk_rlwe;
154 let sk_lwe = &client_key.sk_lwe;
155
156 let d_rgsw_gadget_vec = gadget_vector(
157 self.parameters.rlwe_logq,
158 self.parameters.logb_rgsw,
159 self.parameters.d_rgsw,
160 );
161
162 // generate galois key -g, g
163 let mut galois_keys = HashMap::new();
164 let g = self.parameters.g as isize;
165 for i in [g, -g] {
166 let gk = DefaultSecureRng::with_local_mut(|rng| {
167 let mut ksk_out = M::zeros(self.parameters.d_rgsw * 2, self.parameters.rlwe_n);
168 galois_key_gen(
169 &mut ksk_out,
170 sk_rlwe,
171 i,
172 &d_rgsw_gadget_vec,
173 &self.rlwe_modop,
174 &self.rlwe_nttop,
175 rng,
176 );
177 ksk_out
178 });
179
180 galois_keys.insert(i, gk);
181 }
182
183 // generate rgsw ciphertexts RGSW(si) where si is i^th LWE secret element
184 let ring_size = self.parameters.rlwe_n;
185 let rlwe_q = self.parameters.rlwe_q;
186 let rgsw_cts = sk_lwe
187 .values()
188 .iter()
189 .map(|si| {
190 // X^{si}; assume |emebedding_factor * si| < N
191 let mut m = M::zeros(1, ring_size);
192 let si = (self.embedding_factor as i32) * si;
193 if si < 0 {
194 // X^{-i} = X^{2N - i} = -X^{N-i}
195 m.set(
196 0,
197 ring_size - (si.abs() as usize),
198 rlwe_q - M::MatElement::one(),
199 );
200 } else {
201 // X^{i}
202 m.set(0, (si.abs() as usize), M::MatElement::one());
203 }
204 self.rlwe_nttop.forward(m.get_row_mut(0));
205
206 let rgsw_si = DefaultSecureRng::with_local_mut(|rng| {
207 let mut rgsw_si = M::zeros(self.parameters.d_rgsw * 4, ring_size);
208 encrypt_rgsw(
209 &mut rgsw_si,
210 &m,
211 &d_rgsw_gadget_vec,
212 sk_rlwe,
213 &self.rlwe_modop,
214 &self.rlwe_nttop,
215 rng,
216 );
217 rgsw_si
218 });
219 rgsw_si
220 })
221 .collect_vec();
222
223 // LWE KSK from RLWE secret s -> LWE secret z
224 let d_lwe_gadget = gadget_vector(
225 self.parameters.lwe_logq,
226 self.parameters.logb_lwe,
227 self.parameters.d_lwe,
228 );
229 let mut lwe_ksk = DefaultSecureRng::with_local_mut(|rng| {
230 let mut out = M::zeros(self.parameters.d_lwe * ring_size, self.parameters.lwe_n + 1);
231 lwe_ksk_keygen(
232 &sk_rlwe.values(),
233 &sk_lwe.values(),
234 &mut out,
235 &d_lwe_gadget,
236 &self.lwe_modop,
237 rng,
238 );
239 out
240 });
241
242 ServerKey {
243 rgsw_cts,
244 galois_keys,
245 lwe_ksk,
246 }
247 }
248
249 pub fn encrypt(&self, m: bool, client_key: &ClientKey) -> M::R {
250 let rlwe_q_by8 =
251 M::MatElement::from_f64((self.parameters.rlwe_q.to_f64().unwrap() / 8.0).round())
252 .unwrap();
253 let m = if m {
254 // Q/8
255 rlwe_q_by8
256 } else {
257 // -Q/8
258 self.parameters.rlwe_q - rlwe_q_by8
259 };
260
261 DefaultSecureRng::with_local_mut(|rng| {
262 let mut lwe_out = M::R::zeros(self.parameters.rlwe_n + 1);
263 encrypt_lwe(
264 &mut lwe_out,
265 &m,
266 client_key.sk_rlwe.values(),
267 &self.rlwe_modop,
268 rng,
269 );
270 lwe_out
271 })
272 }
273
274 pub fn decrypt(&self, lwe_ct: &M::R, client_key: &ClientKey) -> bool {
275 let m = decrypt_lwe(lwe_ct, client_key.sk_rlwe.values(), &self.rlwe_modop);
276 let m = {
277 // m + q/8 => {0,q/4 1}
278 let rlwe_q_by8 =
279 M::MatElement::from_f64((self.parameters.rlwe_q.to_f64().unwrap() / 8.0).round())
280 .unwrap();
281 (((m + rlwe_q_by8).to_f64().unwrap() * 4.0) / self.parameters.rlwe_q.to_f64().unwrap())
282 .round()
283 .to_usize()
284 .unwrap()
285 % 4
286 };
287
288 if m == 0 {
289 false
290 } else if m == 1 {
291 true
292 } else {
293 panic!("Incorrect bool decryption. Got m={m} expected m to be 0 or 1")
294 }
295 }
296}
297
298/// LMKCY+ Blind rotation
299///
300/// gk_to_si: [-g^0, -g^1, .., -g^{q/2-1}, g^0, ..., g^{q/2-1}]
301fn blind_rotation<
302 MT: IsTrivial + MatrixMut,
303 Mmut: MatrixMut<MatElement = MT::MatElement> + Matrix,
304 D: Decomposer<Element = MT::MatElement>,
305 NttOp: Ntt<Element = MT::MatElement>,
306 ModOp: ArithmeticOps<Element = MT::MatElement> + VectorOps<Element = MT::MatElement>,
307 K: PbsKey<M = Mmut>,
308>(
309 trivial_rlwe_test_poly: &mut MT,
310 scratch_matrix_dplus2_ring: &mut Mmut,
311 g: isize,
312 w: usize,
313 q: usize,
314 gk_to_si: &[Vec<usize>],
315 decomposer: &D,
316 ntt_op: &NttOp,
317 mod_op: &ModOp,
318 pbs_key: &K,
319) where
320 <Mmut as Matrix>::R: RowMut,
321 Mmut::MatElement: Copy + Zero,
322 <MT as Matrix>::R: RowMut,
323{
324 let q_by_2 = q / 2;
325
326 // -(g^k)
327 for i in 1..q_by_2 {
328 gk_to_si[q_by_2 + i].iter().for_each(|s_index| {
329 rlwe_by_rgsw(
330 trivial_rlwe_test_poly,
331 pbs_key.rgsw_ct_secret_el(*s_index),
332 scratch_matrix_dplus2_ring,
333 decomposer,
334 ntt_op,
335 mod_op,
336 );
337 });
338
339 galois_auto(
340 trivial_rlwe_test_poly,
341 pbs_key.galois_key_for_auto(g),
342 scratch_matrix_dplus2_ring,
343 pbs_key.auto_map_index(g),
344 pbs_key.auto_map_sign(g),
345 mod_op,
346 ntt_op,
347 decomposer,
348 );
349 }
350
351 // -(g^0)
352 gk_to_si[q_by_2].iter().for_each(|s_index| {
353 rlwe_by_rgsw(
354 trivial_rlwe_test_poly,
355 pbs_key.rgsw_ct_secret_el(*s_index),
356 scratch_matrix_dplus2_ring,
357 decomposer,
358 ntt_op,
359 mod_op,
360 );
361 });
362 galois_auto(
363 trivial_rlwe_test_poly,
364 pbs_key.galois_key_for_auto(-g),
365 scratch_matrix_dplus2_ring,
366 pbs_key.auto_map_index(-g),
367 pbs_key.auto_map_sign(-g),
368 mod_op,
369 ntt_op,
370 decomposer,
371 );
372
373 // +(g^k)
374 for i in 1..q_by_2 {
375 gk_to_si[i].iter().for_each(|s_index| {
376 rlwe_by_rgsw(
377 trivial_rlwe_test_poly,
378 pbs_key.rgsw_ct_secret_el(*s_index),
379 scratch_matrix_dplus2_ring,
380 decomposer,
381 ntt_op,
382 mod_op,
383 );
384 });
385
386 galois_auto(
387 trivial_rlwe_test_poly,
388 pbs_key.galois_key_for_auto(g),
389 scratch_matrix_dplus2_ring,
390 pbs_key.auto_map_index(g),
391 pbs_key.auto_map_sign(g),
392 mod_op,
393 ntt_op,
394 decomposer,
395 );
396 }
397
398 // +(g^0)
399 gk_to_si[0].iter().for_each(|s_index| {
400 rlwe_by_rgsw(
401 trivial_rlwe_test_poly,
402 pbs_key.rgsw_ct_secret_el(gk_to_si[q_by_2][*s_index]),
403 scratch_matrix_dplus2_ring,
404 decomposer,
405 ntt_op,
406 mod_op,
407 );
408 });
409}
410
411/// - Mod down
412/// - key switching
413/// - mod down
414/// - blind rotate
415fn pbs<
416 M: Matrix + MatrixMut + MatrixEntity,
417 MT: MatrixMut<MatElement = M::MatElement, R = M::R> + IsTrivial + MatrixEntity,
418 P: Parameters<Element = M::MatElement>,
419 NttOp: Ntt<Element = M::MatElement>,
420 ModOp: ArithmeticOps<Element = M::MatElement> + VectorOps<Element = M::MatElement>,
421 K: PbsKey<M = M>,
422>(
423 parameters: &P,
424 test_vec: &M::R,
425 lwe_in: &mut M::R,
426 lwe_ksk: &M,
427 scratch_lwen_plus1: &mut M::R,
428 scratch_matrix_dplus2_ring: &mut M,
429 modop_lweq: &ModOp,
430 modop_rlweq: &ModOp,
431 nttop_rlweq: &NttOp,
432 pbs_key: K,
433) where
434 <M as Matrix>::R: RowMut,
435 <MT as Matrix>::R: RowMut,
436 M::MatElement: PrimInt + ToPrimitive + FromPrimitive + One + Copy + Zero,
437{
438 let rlwe_q = parameters.rlwe_q();
439 let lwe_q = parameters.lwe_q();
440 let br_q = parameters.br_q();
441 let rlwe_qf64 = rlwe_q.to_f64().unwrap();
442 let lwe_qf64 = lwe_q.to_f64().unwrap();
443 let br_qf64 = br_q.to_f64().unwrap();
444 let rlwe_n = parameters.rlwe_n();
445
446 // moddown Q -> Q_ks
447 lwe_in.as_mut().iter_mut().for_each(|v| {
448 *v =
449 M::MatElement::from_f64(((v.to_f64().unwrap() * lwe_qf64) / rlwe_qf64).round()).unwrap()
450 });
451
452 // key switch
453 // let mut lwe_out = M::zeros(1, parameters.lwe_n() + 1);
454 scratch_lwen_plus1.as_mut().fill(M::MatElement::zero());
455 lwe_key_switch(
456 scratch_lwen_plus1,
457 lwe_in,
458 lwe_ksk,
459 modop_lweq,
460 parameters.decomoposer_lwe(),
461 );
462
463 // odd mowdown Q_ks -> q
464 let g_k_dlog_map = parameters.g_k_dlog_map();
465 let mut g_k_si = vec![vec![]; br_q];
466 scratch_lwen_plus1
467 .as_ref()
468 .iter()
469 .skip(1)
470 .enumerate()
471 .for_each(|(index, v)| {
472 let odd_v = mod_switch_odd(v.to_f64().unwrap(), lwe_qf64, br_qf64);
473 let k = g_k_dlog_map[odd_v];
474 g_k_si[k].push(index);
475 });
476
477 // handle b and set trivial test RLWE
478 let g = parameters.g() as usize;
479 let g_times_b = (g * mod_switch_odd(
480 scratch_lwen_plus1.as_ref()[0].to_f64().unwrap(),
481 lwe_qf64,
482 br_qf64,
483 )) % (br_q);
484 // v = (v(X) * X^{g*b}) mod X^{q/2}+1
485 let br_qby2 = br_q / 2;
486 let mut gb_monomial_sign = true;
487 let mut gb_monomial_exp = g_times_b;
488 // X^{g*b} mod X^{q}+1
489 if gb_monomial_exp > br_qby2 {
490 gb_monomial_exp -= br_qby2;
491 gb_monomial_sign = false
492 }
493 // monomial mul
494 let mut trivial_rlwe_test_poly = MT::zeros(2, rlwe_n);
495 if parameters.embedding_factor() == 1 {
496 monomial_mul(
497 test_vec.as_ref(),
498 trivial_rlwe_test_poly.get_row_mut(1).as_mut(),
499 gb_monomial_exp,
500 gb_monomial_sign,
501 br_q,
502 modop_rlweq,
503 );
504 } else {
505 // use lwe_in to store the `t = v(X) * X^{g*2} mod X^{q/2}+1` temporarily. This
506 // works because q/2 < N (where N is lwe_in LWE dimension) always.
507 monomial_mul(
508 test_vec.as_ref(),
509 &mut lwe_in.as_mut()[..br_qby2],
510 gb_monomial_exp,
511 gb_monomial_sign,
512 br_q,
513 modop_rlweq,
514 );
515
516 // emebed poly `t` in ring X^{q/2}+1 inside the bigger ring X^{N}+1
517 let partb_trivial_rlwe = trivial_rlwe_test_poly.get_row_mut(1);
518 lwe_in.as_ref()[..br_qby2]
519 .iter()
520 .enumerate()
521 .for_each(|(index, v)| {
522 partb_trivial_rlwe[2 * index] = *v;
523 });
524 }
525
526 // blind rotate
527 blind_rotation(
528 &mut trivial_rlwe_test_poly,
529 scratch_matrix_dplus2_ring,
530 parameters.g(),
531 1,
532 br_q,
533 &g_k_si,
534 parameters.decomoposer_rlwe(),
535 nttop_rlweq,
536 modop_rlweq,
537 &pbs_key,
538 );
539
540 // sample extract
541 sample_extract(lwe_in, &trivial_rlwe_test_poly, modop_rlweq, 0);
542}
543
544fn mod_switch_odd(v: f64, from_q: f64, to_q: f64) -> usize {
545 let odd_v = (((v.to_f64().unwrap() * to_q) / (from_q)).floor())
546 .to_usize()
547 .unwrap();
548 //TODO(Jay): check correctness of this
549 odd_v + (odd_v ^ (usize::one()))
550}
551
552fn sample_extract<M: Matrix + MatrixMut, ModOp: ArithmeticOps<Element = M::MatElement>>(
553 lwe_out: &mut M::R,
554 rlwe_in: &M,
555 mod_op: &ModOp,
556 index: usize,
557) where
558 <M as Matrix>::R: RowMut,
559 M::MatElement: Copy,
560{
561 let ring_size = rlwe_in.dimension().1;
562
563 // index..=0
564 let to = &mut lwe_out.as_mut()[1..];
565 let from = rlwe_in.get_row_slice(0);
566 for i in 0..index + 1 {
567 to[i] = from[index - i];
568 }
569
570 // -(N..index)
571 for i in index + 1..ring_size {
572 to[i] = mod_op.neg(&from[ring_size + index - i]);
573 }
574
575 // set b
576 lwe_out.as_mut()[0] = *rlwe_in.get(1, index);
577}
578
579fn monomial_mul<El, ModOp: ArithmeticOps<Element = El>>(
580 p_in: &[El],
581 p_out: &mut [El],
582 mon_exp: usize,
583 mon_sign: bool,
584 ring_size: usize,
585 mod_op: &ModOp,
586) where
587 El: Copy,
588{
589 debug_assert!(p_in.as_ref().len() == ring_size);
590 debug_assert!(p_in.as_ref().len() == p_out.as_ref().len());
591 debug_assert!(mon_exp < ring_size);
592
593 p_in.as_ref().iter().enumerate().for_each(|(index, v)| {
594 let mut to_index = index + mon_exp;
595 let mut to_sign = mon_sign;
596 if to_index >= ring_size {
597 to_index = to_index - ring_size;
598 to_sign = !to_sign;
599 }
600
601 if !to_sign {
602 p_out.as_mut()[to_index] = mod_op.neg(v);
603 } else {
604 p_out.as_mut()[to_index] = *v;
605 }
606 });
607}
608
609#[cfg(test)]
610mod tests {
611 use crate::{backend::ModularOpsU64, ntt::NttBackendU64};
612
613 use super::*;
614
615 const SP_BOOL_PARAMS: BoolParameters<u64> = BoolParameters::<u64> {
616 rlwe_q: 4294957057u64,
617 rlwe_logq: 32,
618 lwe_q: 1 << 16,
619 lwe_logq: 16,
620 br_q: 1 << 9,
621 rlwe_n: 1 << 10,
622 lwe_n: 490,
623 d_rgsw: 4,
624 logb_rgsw: 7,
625 d_lwe: 4,
626 logb_lwe: 4,
627 g: 5,
628 w: 1,
629 };
630
631 #[test]
632 fn encrypt_decrypt_works() {
633 // let prime = generate_prime(32, 2 * 1024, 1 << 32);
634 // dbg!(prime);
635 let bool_evaluator =
636 BoolEvaluator::<Vec<Vec<u64>>, u64, NttBackendU64, ModularOpsU64>::new(SP_BOOL_PARAMS);
637 let client_key = bool_evaluator.client_key();
638 // let sever_key = bool_evaluator.server_key(&client_key);
639
640 let mut m = true;
641 for _ in 0..1000 {
642 let lwe_ct = bool_evaluator.encrypt(m, &client_key);
643 let m_back = bool_evaluator.decrypt(&lwe_ct, &client_key);
644 assert_eq!(m, m_back);
645 m = !m;
646 }
647 }
648}