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