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