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 = M::zeros(2, rlwe_n);
159 if pbs_info.embedding_factor() == 1 {
160 monomial_mul(
161 test_vec.as_ref(),
162 trivial_rlwe_test_poly.get_row_mut(1).as_mut(),
163 gb_monomial_exp,
164 gb_monomial_sign,
165 br_qby2,
166 pbs_info.modop_rlweq(),
167 );
168 } else {
169 // use lwe_in to store the `t = v(X) * X^{g*2} mod X^{q/2}+1` temporarily. This
170 // works because q/2 <= N (where N is lwe_in LWE dimension) always.
171 monomial_mul(
172 test_vec.as_ref(),
173 &mut lwe_in.as_mut()[..br_qby2],
174 gb_monomial_exp,
175 gb_monomial_sign,
176 br_qby2,
177 pbs_info.modop_rlweq(),
178 );
179
180 // emebed poly `t` in ring X^{q/2}+1 inside the bigger ring X^{N}+1
181 let embed_factor = pbs_info.embedding_factor();
182 let partb_trivial_rlwe = trivial_rlwe_test_poly.get_row_mut(1);
183 lwe_in.as_ref()[..br_qby2]
184 .iter()
185 .enumerate()
186 .for_each(|(index, v)| {
187 partb_trivial_rlwe[embed_factor * index] = *v;
188 });
189 }
190
191 // let now = std::time::Instant::now();
192 // blind rotate
193 blind_rotation(
194 &mut trivial_rlwe_test_poly,
195 scratch_blind_rotate_matrix,
196 pbs_info.g(),
197 pbs_info.w(),
198 br_q,
199 &g_k_si,
200 pbs_info.rlwe_rgsw_decomposer(),
201 pbs_info.auto_decomposer(),
202 pbs_info.nttop_rlweq(),
203 pbs_info.modop_rlweq(),
204 pbs_info,
205 pbs_key,
206 );
207 // println!("Blind rotation time: {:?}", now.elapsed());
208
209 // sample extract
210 sample_extract(lwe_in, &trivial_rlwe_test_poly, pbs_info.modop_rlweq(), 0);
211}
212
213/// LMKCY+ Blind rotation
214///
215/// gk_to_si: [g^0, ..., g^{q/2-1}, -g^0, -g^1, .., -g^{q/2-1}]
216fn blind_rotation<
217 Mmut: MatrixMut,
218 D: Decomposer<Element = Mmut::MatElement>,
219 NttOp: Ntt<Element = Mmut::MatElement>,
220 ModOp: ArithmeticOps<Element = Mmut::MatElement> + ShoupMatrixFMA<Mmut::R>,
221 MShoup: WithShoupRepr<M = Mmut>,
222 K: PbsKey<RgswCt = MShoup, AutoKey = MShoup>,
223 P: PbsInfo<M = Mmut>,
224>(
225 trivial_rlwe_test_poly: &mut Mmut,
226 scratch_matrix: &mut Mmut,
227 _g: isize,
228 w: usize,
229 q: usize,
230 gk_to_si: &[Vec<usize>],
231 rlwe_rgsw_decomposer: &(D, D),
232 auto_decomposer: &D,
233 ntt_op: &NttOp,
234 mod_op: &ModOp,
235 parameters: &P,
236 pbs_key: &K,
237) where
238 <Mmut as Matrix>::R: RowMut,
239 Mmut::MatElement: Copy + Zero,
240{
241 let mut is_trivial = true;
242
243 let q_by_4 = q >> 2;
244 let mut count = 0;
245 // -(g^k)
246 let mut v = 0;
247 for i in (1..q_by_4).rev() {
248 // dbg!(q_by_4 + i);
249 let s_indices = &gk_to_si[q_by_4 + i];
250
251 s_indices.iter().for_each(|s_index| {
252 // let new = std::time::Instant::now();
253 let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
254 rlwe_by_rgsw_shoup(
255 trivial_rlwe_test_poly,
256 ct.as_ref(),
257 ct.shoup_repr(),
258 scratch_matrix,
259 rlwe_rgsw_decomposer,
260 ntt_op,
261 mod_op,
262 is_trivial,
263 );
264 is_trivial = false;
265 // println!("Rlwe x Rgsw time: {:?}", new.elapsed());
266 });
267 v += 1;
268
269 if gk_to_si[q_by_4 + i - 1].len() != 0 || v == w || i == 1 {
270 let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(v);
271
272 // let now = std::time::Instant::now();
273 let auto_key = pbs_key.galois_key_for_auto(v);
274 galois_auto_shoup(
275 trivial_rlwe_test_poly,
276 auto_key.as_ref(),
277 auto_key.shoup_repr(),
278 scratch_matrix,
279 &auto_map_index,
280 &auto_map_sign,
281 mod_op,
282 ntt_op,
283 auto_decomposer,
284 is_trivial,
285 );
286 // println!("Auto time: {:?}", now.elapsed());
287
288 count += 1;
289 v = 0;
290 }
291 }
292
293 // -(g^0)
294 {
295 gk_to_si[q_by_4].iter().for_each(|s_index| {
296 let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
297 rlwe_by_rgsw_shoup(
298 trivial_rlwe_test_poly,
299 ct.as_ref(),
300 ct.shoup_repr(),
301 scratch_matrix,
302 rlwe_rgsw_decomposer,
303 ntt_op,
304 mod_op,
305 is_trivial,
306 );
307 is_trivial = false;
308 });
309
310 let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(0);
311 let auto_key = pbs_key.galois_key_for_auto(0);
312 galois_auto_shoup(
313 trivial_rlwe_test_poly,
314 auto_key.as_ref(),
315 auto_key.shoup_repr(),
316 scratch_matrix,
317 &auto_map_index,
318 &auto_map_sign,
319 mod_op,
320 ntt_op,
321 auto_decomposer,
322 is_trivial,
323 );
324 count += 1;
325 }
326
327 // +(g^k)
328 let mut v = 0;
329 for i in (1..q_by_4).rev() {
330 let s_indices = &gk_to_si[i];
331 s_indices.iter().for_each(|s_index| {
332 let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
333 rlwe_by_rgsw_shoup(
334 trivial_rlwe_test_poly,
335 ct.as_ref(),
336 ct.shoup_repr(),
337 scratch_matrix,
338 rlwe_rgsw_decomposer,
339 ntt_op,
340 mod_op,
341 is_trivial,
342 );
343 is_trivial = false;
344 });
345 v += 1;
346
347 if gk_to_si[i - 1].len() != 0 || v == w || i == 1 {
348 let (auto_map_index, auto_map_sign) = parameters.rlwe_auto_map(v);
349 let auto_key = pbs_key.galois_key_for_auto(v);
350 galois_auto_shoup(
351 trivial_rlwe_test_poly,
352 auto_key.as_ref(),
353 auto_key.shoup_repr(),
354 scratch_matrix,
355 &auto_map_index,
356 &auto_map_sign,
357 mod_op,
358 ntt_op,
359 auto_decomposer,
360 is_trivial,
361 );
362 v = 0;
363 count += 1;
364 }
365 }
366
367 // +(g^0)
368 gk_to_si[0].iter().for_each(|s_index| {
369 let ct = pbs_key.rgsw_ct_lwe_si(*s_index);
370 rlwe_by_rgsw_shoup(
371 trivial_rlwe_test_poly,
372 ct.as_ref(),
373 ct.shoup_repr(),
374 scratch_matrix,
375 rlwe_rgsw_decomposer,
376 ntt_op,
377 mod_op,
378 is_trivial,
379 );
380 is_trivial = false;
381 });
382 // println!("Auto count: {count}");
383}
384
385fn mod_switch_odd(v: f64, from_q: f64, to_q: f64) -> usize {
386 let odd_v = (((v * to_q) / (from_q)).floor()).to_usize().unwrap();
387 //TODO(Jay): check correctness of this
388 odd_v + ((odd_v & 1) ^ 1)
389}
390
391// TODO(Jay): Add tests for sample extract
392pub(crate) fn sample_extract<M: Matrix + MatrixMut, ModOp: ArithmeticOps<Element = M::MatElement>>(
393 lwe_out: &mut M::R,
394 rlwe_in: &M,
395 mod_op: &ModOp,
396 index: usize,
397) where
398 <M as Matrix>::R: RowMut,
399 M::MatElement: Copy,
400{
401 let ring_size = rlwe_in.dimension().1;
402 assert!(ring_size + 1 == lwe_out.as_ref().len());
403
404 // index..=0
405 let to = &mut lwe_out.as_mut()[1..];
406 let from = rlwe_in.get_row_slice(0);
407 for i in 0..index + 1 {
408 to[i] = from[index - i];
409 }
410
411 // -(N..index)
412 for i in index + 1..ring_size {
413 to[i] = mod_op.neg(&from[ring_size + index - i]);
414 }
415
416 // set b
417 lwe_out.as_mut()[0] = *rlwe_in.get(1, index);
418}
419
420/// Monomial multiplication (p(X)*X^{mon_exp})
421///
422/// - p_out: Output is written to p_out and independent of values in p_out
423fn monomial_mul<El, ModOp: ArithmeticOps<Element = El>>(
424 p_in: &[El],
425 p_out: &mut [El],
426 mon_exp: usize,
427 mon_sign: bool,
428 ring_size: usize,
429 mod_op: &ModOp,
430) where
431 El: Copy,
432{
433 debug_assert!(p_in.as_ref().len() == ring_size);
434 debug_assert!(p_in.as_ref().len() == p_out.as_ref().len());
435 debug_assert!(mon_exp < ring_size);
436
437 p_in.as_ref().iter().enumerate().for_each(|(index, v)| {
438 let mut to_index = index + mon_exp;
439 let mut to_sign = mon_sign;
440 if to_index >= ring_size {
441 to_index = to_index - ring_size;
442 to_sign = !to_sign;
443 }
444
445 if !to_sign {
446 p_out.as_mut()[to_index] = mod_op.neg(v);
447 } else {
448 p_out.as_mut()[to_index] = *v;
449 }
450 });
451}