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