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