phantom-zone/src/pbs.rs
src/pbs.rsRaw
1use std::fmt::Display;
2
3use num_traits::{FromPrimitive, One, PrimInt, ToPrimitive, Zero};
4
5use crate::{
6 backend::{ArithmeticOps, Modulus, ShoupMatrixFMA, VectorOps},
7 decomposer::{Decomposer, RlweDecomposer},
8 lwe::lwe_key_switch,
9 ntt::Ntt,
10 rgsw::{
11 rlwe_auto_shoup, rlwe_by_rgsw_shoup, RgswCiphertextRef, RlweCiphertextMutRef, RlweKskRef,
12 RuntimeScratchMutRef,
13 },
14 Matrix, MatrixEntity, MatrixMut, RowMut,
15};
16pub(crate) trait PbsKey {
17 type RgswCt;
18 type AutoKey;
19 type LweKskKey;
20
21 /// RGSW ciphertext of LWE secret elements
22 fn rgsw_ct_lwe_si(&self, si: usize) -> &Self::RgswCt;
23 /// Key for automorphism with g^k. For -g use k = 0
24 fn galois_key_for_auto(&self, k: usize) -> &Self::AutoKey;
25 /// LWE ksk to key switch from RLWE secret to LWE secret
26 fn lwe_ksk(&self) -> &Self::LweKskKey;
27}
28
29pub(crate) trait WithShoupRepr: AsRef<Self::M> {
30 type M;
31 fn shoup_repr(&self) -> &Self::M;
32}
33
34pub(crate) trait PbsInfo {
35 /// Type of Matrix
36 type M: Matrix;
37 /// Type of Ciphertext modulus
38 type Modulus: Modulus<Element = <Self::M as Matrix>::MatElement>;
39 /// Type of Ntt Operator for Ring polynomials
40 type NttOp: Ntt<Element = <Self::M as Matrix>::MatElement>;
41 /// Type of Signed Decomposer
42 type D: Decomposer<Element = <Self::M as Matrix>::MatElement>;
43
44 // Although both `RlweModOp` and `LweModOp` types have same bounds, they can be
45 // different types. For ex, type RlweModOp may only support native modulus,
46 // where LweModOp may only support prime modulus, etc.
47
48 /// Type of RLWE Modulus Operator
49 type RlweModOp: ArithmeticOps<Element = <Self::M as Matrix>::MatElement>
50 + ShoupMatrixFMA<<Self::M as Matrix>::R>;
51 /// Type of LWE Modulus Operator
52 type LweModOp: VectorOps<Element = <Self::M as Matrix>::MatElement>
53 + ArithmeticOps<Element = <Self::M as Matrix>::MatElement>;
54
55 /// RLWE ciphertext modulus
56 fn rlwe_q(&self) -> &Self::Modulus;
57 /// LWE ciphertext modulus
58 fn lwe_q(&self) -> &Self::Modulus;
59 /// Blind rotation modulus. It is the modulus to which we switch for blind
60 /// rotation. Since blind rotation decrypts LWE ciphetext in the exponent of
61 /// ring polynmial (which is a ring mod 2N), `br_q <= 2N`
62 fn br_q(&self) -> usize;
63 /// Ring polynomial size `N`
64 fn rlwe_n(&self) -> usize;
65 /// LWE dimension `n`
66 fn lwe_n(&self) -> usize;
67 /// Embedding fator for ring X^{q}+1 inside
68 fn embedding_factor(&self) -> usize;
69 /// Window size parameter LKMC++ blind rotaiton
70 fn w(&self) -> usize;
71 /// generator `g` for group Z^*_{br_q}
72 fn g(&self) -> isize;
73 /// LWE key switching decomposer
74 fn lwe_decomposer(&self) -> &Self::D;
75 /// RLWE x RGSW decoposer
76 fn rlwe_rgsw_decomposer(&self) -> &(Self::D, Self::D);
77 /// RLWE auto decomposer
78 fn auto_decomposer(&self) -> &Self::D;
79
80 /// LWE modulus operator
81 fn modop_lweq(&self) -> &Self::LweModOp;
82 /// RLWE modulus operator
83 fn modop_rlweq(&self) -> &Self::RlweModOp;
84
85 /// Ntt operators
86 fn nttop_rlweq(&self) -> &Self::NttOp;
87
88 /// Maps a \in Z^*_{br_q} to discrete log k, with generator g (i.e. g^k =
89 /// a). Returned vector is of size q that stores dlog of `a` at `vec[a]`.
90 ///
91 /// For any `a`, if k is s.t. `a = g^{k} % br_q`, then `k` is expressed as
92 /// k. If `k` is s.t `a = -g^{k} % br_q`, then `k` is expressed as
93 /// k=k+q/4
94 fn g_k_dlog_map(&self) -> &[usize];
95 /// Returns auto map and index vector for auto element g^k. For auto element
96 /// -g set k = 0.
97 fn rlwe_auto_map(&self, k: usize) -> &(Vec<usize>, Vec<bool>);
98}
99
100/// - Mod down
101/// - key switching
102/// - mod down
103/// - blind rotate
104pub(crate) fn pbs<
105 M: MatrixMut + MatrixEntity,
106 MShoup: WithShoupRepr<M = M>,
107 P: PbsInfo<M = M>,
108 K: PbsKey<RgswCt = MShoup, AutoKey = MShoup, LweKskKey = M>,
109>(
110 pbs_info: &P,
111 test_vec: &M::R,
112 lwe_in: &mut M::R,
113 pbs_key: &K,
114 scratch_lwe_vec: &mut M::R,
115 scratch_blind_rotate_matrix: &mut M,
116) where
117 <M as Matrix>::R: RowMut,
118 M::MatElement: PrimInt + FromPrimitive + One + Copy + Zero + Display,
119{
120 let rlwe_q = pbs_info.rlwe_q();
121 let lwe_q = pbs_info.lwe_q();
122 let br_q = pbs_info.br_q();
123 let rlwe_qf64 = rlwe_q.q_as_f64().unwrap();
124 let lwe_qf64 = lwe_q.q_as_f64().unwrap();
125 let br_qf64 = br_q.to_f64().unwrap();
126 let rlwe_n = pbs_info.rlwe_n();
127
128 // moddown Q -> Q_ks
129 lwe_in.as_mut().iter_mut().for_each(|v| {
130 *v =
131 M::MatElement::from_f64(((v.to_f64().unwrap() * lwe_qf64) / rlwe_qf64).round()).unwrap()
132 });
133
134 // key switch RLWE secret to LWE secret
135 // let now = std::time::Instant::now();
136 scratch_lwe_vec.as_mut().fill(M::MatElement::zero());
137 lwe_key_switch(
138 scratch_lwe_vec,
139 lwe_in,
140 pbs_key.lwe_ksk(),
141 pbs_info.modop_lweq(),
142 pbs_info.lwe_decomposer(),
143 );
144 // println!("Time: {:?}", now.elapsed());
145
146 // odd moddown Q_ks -> q
147 let g_k_dlog_map = pbs_info.g_k_dlog_map();
148 let mut g_k_si = vec![vec![]; br_q >> 1];
149 scratch_lwe_vec
150 .as_ref()
151 .iter()
152 .skip(1)
153 .enumerate()
154 .for_each(|(index, v)| {
155 let odd_v = mod_switch_odd(v.to_f64().unwrap(), lwe_qf64, br_qf64);
156 // dlog `k` for `odd_v` is stored as `k` if odd_v = +g^{k}. If odd_v = -g^{k},
157 // then `k` is stored as `q/4 + k`.
158 let k = g_k_dlog_map[odd_v];
159 // assert!(k != 0);
160 g_k_si[k].push(index);
161 });
162
163 // handle b and set trivial test RLWE
164 let g = pbs_info.g() as usize;
165 let g_times_b = (g * mod_switch_odd(
166 scratch_lwe_vec.as_ref()[0].to_f64().unwrap(),
167 lwe_qf64,
168 br_qf64,
169 )) % (br_q);
170 // v = (v(X) * X^{g*b}) mod X^{q/2}+1
171 let br_qby2 = br_q >> 1;
172 let mut gb_monomial_sign = true;
173 let mut gb_monomial_exp = g_times_b;
174 // X^{g*b} mod X^{q/2}+1
175 if gb_monomial_exp > br_qby2 {
176 gb_monomial_exp -= br_qby2;
177 gb_monomial_sign = false
178 }
179 // monomial mul
180 let mut trivial_rlwe_test_poly = M::zeros(2, rlwe_n);
181 if pbs_info.embedding_factor() == 1 {
182 monomial_mul(
183 test_vec.as_ref(),
184 trivial_rlwe_test_poly.get_row_mut(1).as_mut(),
185 gb_monomial_exp,
186 gb_monomial_sign,
187 br_qby2,
188 pbs_info.modop_rlweq(),
189 );
190 } else {
191 // use lwe_in to store the `t = v(X) * X^{g*2} mod X^{q/2}+1` temporarily. This
192 // works because q/2 <= N (where N is lwe_in LWE dimension) always.
193 monomial_mul(
194 test_vec.as_ref(),
195 &mut lwe_in.as_mut()[..br_qby2],
196 gb_monomial_exp,
197 gb_monomial_sign,
198 br_qby2,
199 pbs_info.modop_rlweq(),
200 );
201
202 // emebed poly `t` in ring X^{q/2}+1 inside the bigger ring X^{N}+1
203 let embed_factor = pbs_info.embedding_factor();
204 let partb_trivial_rlwe = trivial_rlwe_test_poly.get_row_mut(1);
205 lwe_in.as_ref()[..br_qby2]
206 .iter()
207 .enumerate()
208 .for_each(|(index, v)| {
209 partb_trivial_rlwe[embed_factor * index] = *v;
210 });
211 }
212
213 // let now = std::time::Instant::now();
214 // blind rotate
215 blind_rotation(
216 &mut trivial_rlwe_test_poly,
217 scratch_blind_rotate_matrix,
218 pbs_info.g(),
219 pbs_info.w(),
220 br_q,
221 &g_k_si,
222 pbs_info.rlwe_rgsw_decomposer(),
223 pbs_info.auto_decomposer(),
224 pbs_info.nttop_rlweq(),
225 pbs_info.modop_rlweq(),
226 pbs_info,
227 pbs_key,
228 );
229 // println!("Blind rotation time: {:?}", now.elapsed());
230
231 // sample extract
232 sample_extract(lwe_in, &trivial_rlwe_test_poly, pbs_info.modop_rlweq(), 0);
233}
234
235/// LMKCY+ Blind rotation
236///
237/// - gk_to_si: Contains LWE secret index `i` in array of secret indices at k^th
238/// index if a_i = g^k if k < q/4 or a_i = -g^k if k > q/4. [g^0, ...,
239/// g^{q/2-1}, -g^0, -g^1, .., -g^{q/2-1}]
240fn blind_rotation<
241 Mmut: MatrixMut,
242 RlweD: RlweDecomposer<Element = Mmut::MatElement>,
243 AutoD: Decomposer<Element = Mmut::MatElement>,
244 NttOp: Ntt<Element = Mmut::MatElement>,
245 ModOp: ArithmeticOps<Element = Mmut::MatElement> + ShoupMatrixFMA<Mmut::R>,
246 MShoup: WithShoupRepr<M = Mmut>,
247 K: PbsKey<RgswCt = MShoup, AutoKey = MShoup>,
248 P: PbsInfo<M = Mmut>,
249>(
250 trivial_rlwe_test_poly: &mut Mmut,
251 scratch_matrix: &mut Mmut,
252 _g: isize,
253 w: usize,
254 q: usize,
255 gk_to_si: &[Vec<usize>],
256 rlwe_rgsw_decomposer: &RlweD,
257 auto_decomposer: &AutoD,
258 ntt_op: &NttOp,
259 mod_op: &ModOp,
260 parameters: &P,
261 pbs_key: &K,
262) where
263 <Mmut as Matrix>::R: RowMut,
264 Mmut::MatElement: Copy + Zero,
265{
266 let mut is_trivial = true;
267 let mut scratch_matrix = RuntimeScratchMutRef::new(scratch_matrix.as_mut());
268 let mut rlwe = RlweCiphertextMutRef::new(trivial_rlwe_test_poly.as_mut());
269 let d_a = rlwe_rgsw_decomposer.a().decomposition_count().0;
270 let d_b = rlwe_rgsw_decomposer.b().decomposition_count().0;
271 let d_auto = auto_decomposer.decomposition_count().0;
272
273 let q_by_4 = q >> 2;
274 // let mut count = 0;
275 // -(g^k)
276 let mut v = 0;
277 for i in (1..q_by_4).rev() {
278 // dbg!(q_by_4 + i);
279 let s_indices = &gk_to_si[q_by_4 + i];
280
281 s_indices.iter().for_each(|s_index| {
282 // let new = std::time::Instant::now();
283 let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
284 rlwe_by_rgsw_shoup(
285 &mut rlwe,
286 &RgswCiphertextRef::new(ct.as_ref().as_ref(), d_a, d_b),
287 &RgswCiphertextRef::new(ct.shoup_repr().as_ref(), d_a, d_b),
288 &mut scratch_matrix,
289 rlwe_rgsw_decomposer,
290 ntt_op,
291 mod_op,
292 is_trivial,
293 );
294 is_trivial = false;
295 // println!("Rlwe x Rgsw time: {:?}", new.elapsed());
296 });
297 v += 1;
298
299 if gk_to_si[q_by_4 + i - 1].len() != 0 || v == w || i == 1 {
300 let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(v);
301
302 // let now = std::time::Instant::now();
303 let auto_key = pbs_key.galois_key_for_auto(v);
304 rlwe_auto_shoup(
305 &mut rlwe,
306 &RlweKskRef::new(auto_key.as_ref().as_ref(), d_auto),
307 &RlweKskRef::new(auto_key.shoup_repr().as_ref(), d_auto),
308 &mut scratch_matrix,
309 &auto_map_index,
310 &auto_map_sign,
311 mod_op,
312 ntt_op,
313 auto_decomposer,
314 is_trivial,
315 );
316 // println!("Auto time: {:?}", now.elapsed());
317 // count += 1;
318
319 v = 0;
320 }
321 }
322
323 // -(g^0)
324 {
325 gk_to_si[q_by_4].iter().for_each(|s_index| {
326 let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
327 rlwe_by_rgsw_shoup(
328 &mut rlwe,
329 &RgswCiphertextRef::new(ct.as_ref().as_ref(), d_a, d_b),
330 &RgswCiphertextRef::new(ct.shoup_repr().as_ref(), d_a, d_b),
331 &mut scratch_matrix,
332 rlwe_rgsw_decomposer,
333 ntt_op,
334 mod_op,
335 is_trivial,
336 );
337 is_trivial = false;
338 });
339
340 let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(0);
341 let auto_key = pbs_key.galois_key_for_auto(0);
342 rlwe_auto_shoup(
343 &mut rlwe,
344 &RlweKskRef::new(auto_key.as_ref().as_ref(), d_auto),
345 &RlweKskRef::new(auto_key.shoup_repr().as_ref(), d_auto),
346 &mut scratch_matrix,
347 &auto_map_index,
348 &auto_map_sign,
349 mod_op,
350 ntt_op,
351 auto_decomposer,
352 is_trivial,
353 );
354 // count += 1;
355 }
356
357 // +(g^k)
358 let mut v = 0;
359 for i in (1..q_by_4).rev() {
360 let s_indices = &gk_to_si[i];
361 s_indices.iter().for_each(|s_index| {
362 let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
363 rlwe_by_rgsw_shoup(
364 &mut rlwe,
365 &RgswCiphertextRef::new(ct.as_ref().as_ref(), d_a, d_b),
366 &RgswCiphertextRef::new(ct.shoup_repr().as_ref(), d_a, d_b),
367 &mut scratch_matrix,
368 rlwe_rgsw_decomposer,
369 ntt_op,
370 mod_op,
371 is_trivial,
372 );
373 is_trivial = false;
374 });
375 v += 1;
376
377 if gk_to_si[i - 1].len() != 0 || v == w || i == 1 {
378 let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(v);
379 let auto_key = pbs_key.galois_key_for_auto(v);
380 rlwe_auto_shoup(
381 &mut rlwe,
382 &RlweKskRef::new(auto_key.as_ref().as_ref(), d_auto),
383 &RlweKskRef::new(auto_key.shoup_repr().as_ref(), d_auto),
384 &mut scratch_matrix,
385 &auto_map_index,
386 &auto_map_sign,
387 mod_op,
388 ntt_op,
389 auto_decomposer,
390 is_trivial,
391 );
392 v = 0;
393
394 // count += 1;
395 }
396 }
397
398 // +(g^0)
399 gk_to_si[0].iter().for_each(|s_index| {
400 let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
401 rlwe_by_rgsw_shoup(
402 &mut rlwe,
403 &RgswCiphertextRef::new(ct.as_ref().as_ref(), d_a, d_b),
404 &RgswCiphertextRef::new(ct.shoup_repr().as_ref(), d_a, d_b),
405 &mut scratch_matrix,
406 rlwe_rgsw_decomposer,
407 ntt_op,
408 mod_op,
409 is_trivial,
410 );
411 is_trivial = false;
412 });
413 // println!("Auto count: {count}");
414}
415
416fn mod_switch_odd(v: f64, from_q: f64, to_q: f64) -> usize {
417 let odd_v = (((v * to_q) / (from_q)).floor()).to_usize().unwrap();
418 //TODO(Jay): check correctness of this
419 odd_v + ((odd_v & 1) ^ 1)
420}
421
422// TODO(Jay): Add tests for sample extract
423pub(crate) fn sample_extract<M: Matrix + MatrixMut, ModOp: ArithmeticOps<Element = M::MatElement>>(
424 lwe_out: &mut M::R,
425 rlwe_in: &M,
426 mod_op: &ModOp,
427 index: usize,
428) where
429 <M as Matrix>::R: RowMut,
430 M::MatElement: Copy,
431{
432 let ring_size = rlwe_in.dimension().1;
433 assert!(ring_size + 1 == lwe_out.as_ref().len());
434
435 // index..=0
436 let to = &mut lwe_out.as_mut()[1..];
437 let from = rlwe_in.get_row_slice(0);
438 for i in 0..index + 1 {
439 to[i] = from[index - i];
440 }
441
442 // -(N..index)
443 for i in index + 1..ring_size {
444 to[i] = mod_op.neg(&from[ring_size + index - i]);
445 }
446
447 // set b
448 lwe_out.as_mut()[0] = *rlwe_in.get(1, index);
449}
450
451/// Monomial multiplication (p(X)*X^{mon_exp})
452///
453/// - p_out: Output is written to p_out and independent of values in p_out
454fn monomial_mul<El, ModOp: ArithmeticOps<Element = El>>(
455 p_in: &[El],
456 p_out: &mut [El],
457 mon_exp: usize,
458 mon_sign: bool,
459 ring_size: usize,
460 mod_op: &ModOp,
461) where
462 El: Copy,
463{
464 debug_assert!(p_in.as_ref().len() == ring_size);
465 debug_assert!(p_in.as_ref().len() == p_out.as_ref().len());
466 debug_assert!(mon_exp < ring_size);
467
468 p_in.as_ref().iter().enumerate().for_each(|(index, v)| {
469 let mut to_index = index + mon_exp;
470 let mut to_sign = mon_sign;
471 if to_index >= ring_size {
472 to_index = to_index - ring_size;
473 to_sign = !to_sign;
474 }
475
476 if !to_sign {
477 p_out.as_mut()[to_index] = mod_op.neg(v);
478 } else {
479 p_out.as_mut()[to_index] = *v;
480 }
481 });
482}