phantom-zone/src/lwe.rs
src/lwe.rsRaw
1use std::{
2 cell::RefCell,
3 collections::btree_map::Values,
4 fmt::{Debug, Display},
5 marker::PhantomData,
6};
7
8use itertools::{izip, Itertools};
9use num_traits::{abs, PrimInt, ToPrimitive, Zero};
10
11use crate::{
12 backend::{ArithmeticOps, GetModulus, Modulus, VectorOps},
13 decomposer::Decomposer,
14 random::{
15 DefaultSecureRng, NewWithSeed, RandomFillGaussianInModulus, RandomFillUniformInModulus,
16 RandomGaussianElementInModulus, DEFAULT_RNG,
17 },
18 utils::{fill_random_ternary_secret_with_hamming_weight, TryConvertFrom1, WithLocal},
19 Matrix, MatrixEntity, MatrixMut, Row, RowEntity, RowMut, Secret,
20};
21
22struct SeededLweKeySwitchingKey<Ro, S>
23where
24 Ro: Row,
25{
26 data: Ro,
27 seed: S,
28 to_lwe_n: usize,
29 modulus: Ro::Element,
30}
31
32impl<Ro: RowEntity, S> SeededLweKeySwitchingKey<Ro, S> {
33 pub(crate) fn empty(
34 from_lwe_n: usize,
35 to_lwe_n: usize,
36 d: usize,
37 seed: S,
38 modulus: Ro::Element,
39 ) -> Self {
40 let data = Ro::zeros(from_lwe_n * d);
41 SeededLweKeySwitchingKey {
42 data,
43 to_lwe_n,
44 seed,
45 modulus,
46 }
47 }
48}
49
50struct LweKeySwitchingKey<M, R> {
51 data: M,
52 _phantom: PhantomData<R>,
53}
54
55impl<
56 M: MatrixMut + MatrixEntity,
57 R: NewWithSeed + RandomFillUniformInModulus<[M::MatElement], M::MatElement>,
58 > From<&SeededLweKeySwitchingKey<M::R, R::Seed>> for LweKeySwitchingKey<M, R>
59where
60 M::R: RowMut,
61 R::Seed: Clone,
62 M::MatElement: Copy,
63{
64 fn from(value: &SeededLweKeySwitchingKey<M::R, R::Seed>) -> Self {
65 let mut p_rng = R::new_with_seed(value.seed.clone());
66 let mut data = M::zeros(value.data.as_ref().len(), value.to_lwe_n + 1);
67 izip!(value.data.as_ref().iter(), data.iter_rows_mut()).for_each(|(bi, lwe_i)| {
68 RandomFillUniformInModulus::random_fill(
69 &mut p_rng,
70 &value.modulus,
71 &mut lwe_i.as_mut()[1..],
72 );
73 lwe_i.as_mut()[0] = *bi;
74 });
75 LweKeySwitchingKey {
76 data,
77 _phantom: PhantomData,
78 }
79 }
80}
81
82trait LweCiphertext<M: Matrix> {}
83
84#[derive(Clone)]
85pub struct LweSecret {
86 pub(crate) values: Vec<i32>,
87}
88
89impl Secret for LweSecret {
90 type Element = i32;
91 fn values(&self) -> &[Self::Element] {
92 &self.values
93 }
94}
95
96impl LweSecret {
97 pub(crate) fn random(hw: usize, n: usize) -> LweSecret {
98 DefaultSecureRng::with_local_mut(|rng| {
99 let mut out = vec![0i32; n];
100 fill_random_ternary_secret_with_hamming_weight(&mut out, hw, rng);
101
102 LweSecret { values: out }
103 })
104 }
105}
106
107pub(crate) fn lwe_key_switch<
108 M: Matrix,
109 Ro: AsMut<[M::MatElement]> + AsRef<[M::MatElement]>,
110 Op: VectorOps<Element = M::MatElement> + ArithmeticOps<Element = M::MatElement>,
111 D: Decomposer<Element = M::MatElement>,
112>(
113 lwe_out: &mut Ro,
114 lwe_in: &Ro,
115 lwe_ksk: &M,
116 operator: &Op,
117 decomposer: &D,
118) {
119 assert!(
120 lwe_ksk.dimension().0 == ((lwe_in.as_ref().len() - 1) * decomposer.decomposition_count())
121 );
122 assert!(lwe_out.as_ref().len() == lwe_ksk.dimension().1);
123
124 let lwe_in_a_decomposed = lwe_in
125 .as_ref()
126 .iter()
127 .skip(1)
128 .flat_map(|ai| decomposer.decompose(ai));
129 izip!(lwe_in_a_decomposed, lwe_ksk.iter_rows()).for_each(|(ai_j, beta_ij_lwe)| {
130 operator.elwise_fma_scalar_mut(lwe_out.as_mut(), beta_ij_lwe.as_ref(), &ai_j);
131 });
132
133 let out_b = operator.add(&lwe_out.as_ref()[0], &lwe_in.as_ref()[0]);
134 lwe_out.as_mut()[0] = out_b;
135}
136
137pub fn lwe_ksk_keygen<
138 Ro: Row + RowMut + RowEntity,
139 S,
140 Op: VectorOps<Element = Ro::Element>
141 + ArithmeticOps<Element = Ro::Element>
142 + GetModulus<Element = Ro::Element>,
143 R: RandomGaussianElementInModulus<Ro::Element, Op::M>,
144 PR: RandomFillUniformInModulus<[Ro::Element], Op::M>,
145>(
146 from_lwe_sk: &[S],
147 to_lwe_sk: &[S],
148 ksk_out: &mut Ro,
149 gadget: &[Ro::Element],
150 operator: &Op,
151 p_rng: &mut PR,
152 rng: &mut R,
153) where
154 Ro: TryConvertFrom1<[S], Op::M>,
155 Ro::Element: Zero + Debug,
156{
157 assert!(ksk_out.as_ref().len() == (from_lwe_sk.len() * gadget.len()));
158
159 let d = gadget.len();
160
161 let modulus = operator.modulus();
162 let mut neg_sk_in_m = Ro::try_convert_from(from_lwe_sk, modulus);
163 operator.elwise_neg_mut(neg_sk_in_m.as_mut());
164 let sk_out_m = Ro::try_convert_from(to_lwe_sk, modulus);
165
166 let mut scratch = Ro::zeros(to_lwe_sk.len());
167
168 izip!(neg_sk_in_m.as_ref(), ksk_out.as_mut().chunks_mut(d)).for_each(
169 |(neg_sk_in_si, d_lwes_partb)| {
170 izip!(gadget.iter(), d_lwes_partb.into_iter()).for_each(|(f, lwe_b)| {
171 // sample `a`
172 RandomFillUniformInModulus::random_fill(p_rng, &modulus, scratch.as_mut());
173
174 // a * z
175 let mut az = Ro::Element::zero();
176 izip!(scratch.as_ref().iter(), sk_out_m.as_ref()).for_each(|(ai, si)| {
177 let ai_si = operator.mul(ai, si);
178 az = operator.add(&az, &ai_si);
179 });
180
181 // a*z + (-s_i)*\beta^j + e
182 let mut b = operator.add(&az, &operator.mul(f, neg_sk_in_si));
183 let e = RandomGaussianElementInModulus::random(rng, &modulus);
184 b = operator.add(&b, &e);
185
186 *lwe_b = b;
187 })
188 },
189 );
190}
191
192/// Encrypts encoded message m as LWE ciphertext
193pub fn encrypt_lwe<
194 Ro: Row + RowMut,
195 Op: ArithmeticOps<Element = Ro::Element> + GetModulus<Element = Ro::Element>,
196 R: RandomGaussianElementInModulus<Ro::Element, Op::M>
197 + RandomFillUniformInModulus<[Ro::Element], Op::M>,
198 S,
199>(
200 lwe_out: &mut Ro,
201 m: &Ro::Element,
202 s: &[S],
203 operator: &Op,
204 rng: &mut R,
205) where
206 Ro: TryConvertFrom1<[S], Op::M>,
207 Ro::Element: Zero,
208{
209 let s = Ro::try_convert_from(s, operator.modulus());
210 assert!(s.as_ref().len() == (lwe_out.as_ref().len() - 1));
211
212 // a*s
213 RandomFillUniformInModulus::random_fill(rng, operator.modulus(), &mut lwe_out.as_mut()[1..]);
214 let mut sa = Ro::Element::zero();
215 izip!(lwe_out.as_mut().iter().skip(1), s.as_ref()).for_each(|(ai, si)| {
216 let tmp = operator.mul(ai, si);
217 sa = operator.add(&tmp, &sa);
218 });
219
220 // b = a*s + e + m
221 let e = RandomGaussianElementInModulus::random(rng, operator.modulus());
222 let b = operator.add(&operator.add(&sa, &e), m);
223 lwe_out.as_mut()[0] = b;
224}
225
226pub fn decrypt_lwe<
227 Ro: Row,
228 Op: ArithmeticOps<Element = Ro::Element> + GetModulus<Element = Ro::Element>,
229 S,
230>(
231 lwe_ct: &Ro,
232 s: &[S],
233 operator: &Op,
234) -> Ro::Element
235where
236 Ro: TryConvertFrom1<[S], Op::M>,
237 Ro::Element: Zero,
238{
239 let s = Ro::try_convert_from(s, operator.modulus());
240
241 let mut sa = Ro::Element::zero();
242 izip!(lwe_ct.as_ref().iter().skip(1), s.as_ref()).for_each(|(ai, si)| {
243 let tmp = operator.mul(ai, si);
244 sa = operator.add(&tmp, &sa);
245 });
246
247 let b = &lwe_ct.as_ref()[0];
248 operator.sub(b, &sa)
249}
250
251/// Measures noise in input LWE ciphertext with reference of `ideal_m`
252///
253/// - ct: Input LWE ciphertext
254/// - s: corresponding secret
255/// - ideal_m: Ideal `encoded` message
256pub(crate) fn measure_noise_lwe<
257 Ro: Row,
258 Op: ArithmeticOps<Element = Ro::Element> + GetModulus<Element = Ro::Element>,
259 S,
260>(
261 ct: &Ro,
262 s: &[S],
263 operator: &Op,
264 ideal_m: &Ro::Element,
265) -> f64
266where
267 Ro: TryConvertFrom1<[S], Op::M>,
268 Ro::Element: Zero + ToPrimitive + PrimInt + Display,
269{
270 assert!(s.len() == ct.as_ref().len() - 1,);
271
272 let s = Ro::try_convert_from(s, &operator.modulus());
273 let mut sa = Ro::Element::zero();
274 izip!(s.as_ref().iter(), ct.as_ref().iter().skip(1)).for_each(|(si, ai)| {
275 sa = operator.add(&sa, &operator.mul(si, ai));
276 });
277 let m = operator.sub(&ct.as_ref()[0], &sa);
278
279 let mut diff = operator.sub(&m, ideal_m);
280 let q = operator.modulus();
281 return q.map_element_to_i64(&diff).to_f64().unwrap().abs().log2();
282}
283
284#[cfg(test)]
285mod tests {
286
287 use crate::{
288 backend::{ModInit, ModularOpsU64},
289 decomposer::{Decomposer, DefaultDecomposer},
290 lwe::{lwe_key_switch, measure_noise_lwe},
291 random::DefaultSecureRng,
292 rgsw::measure_noise,
293 Secret,
294 };
295
296 use super::{
297 decrypt_lwe, encrypt_lwe, lwe_ksk_keygen, LweKeySwitchingKey, LweSecret,
298 SeededLweKeySwitchingKey,
299 };
300
301 const K: usize = 50;
302
303 #[test]
304 fn encrypt_decrypt_works() {
305 let logq = 16;
306 let q = 1u64 << logq;
307 let lwe_n = 1024;
308 let logp = 3;
309
310 let modq_op = ModularOpsU64::new(q);
311 let lwe_sk = LweSecret::random(lwe_n >> 1, lwe_n);
312
313 let mut rng = DefaultSecureRng::new();
314
315 // encrypt
316 for m in 0..1u64 << logp {
317 let encoded_m = m << (logq - logp);
318 let mut lwe_ct = vec![0u64; lwe_n + 1];
319 encrypt_lwe(
320 &mut lwe_ct,
321 &encoded_m,
322 &lwe_sk.values(),
323 &modq_op,
324 &mut rng,
325 );
326 let encoded_m_back = decrypt_lwe(&lwe_ct, &lwe_sk.values(), &modq_op);
327 let m_back = ((((encoded_m_back as f64) * ((1 << logp) as f64)) / q as f64).round()
328 as u64)
329 % (1u64 << logp);
330 assert_eq!(m, m_back, "Expected {m} but got {m_back}");
331 }
332 }
333
334 #[test]
335 fn key_switch_works() {
336 let logq = 18;
337 let logp = 2;
338 let q = 1u64 << logq;
339 let lwe_in_n = 2048;
340 let lwe_out_n = 493;
341 let d_ks = 3;
342 let logb = 6;
343
344 let lwe_sk_in = LweSecret::random(lwe_in_n >> 1, lwe_in_n);
345 let lwe_sk_out = LweSecret::random(lwe_out_n >> 1, lwe_out_n);
346
347 let mut rng = DefaultSecureRng::new();
348 let modq_op = ModularOpsU64::new(q);
349
350 // genrate ksk
351 for _ in 0..K {
352 let mut ksk_seed = [0u8; 32];
353 rng.fill_bytes(&mut ksk_seed);
354 let mut seeded_ksk =
355 SeededLweKeySwitchingKey::empty(lwe_in_n, lwe_out_n, d_ks, ksk_seed, q);
356 let mut p_rng = DefaultSecureRng::new_seeded(ksk_seed);
357 let decomposer = DefaultDecomposer::new(q, logb, d_ks);
358 let gadget = decomposer.gadget_vector();
359 lwe_ksk_keygen(
360 &lwe_sk_in.values(),
361 &lwe_sk_out.values(),
362 &mut seeded_ksk.data,
363 &gadget,
364 &modq_op,
365 &mut p_rng,
366 &mut rng,
367 );
368 // println!("{:?}", ksk);
369 let ksk = LweKeySwitchingKey::<Vec<Vec<u64>>, DefaultSecureRng>::from(&seeded_ksk);
370
371 for m in 0..(1 << logp) {
372 // encrypt using lwe_sk_in
373 let encoded_m = m << (logq - logp);
374 let mut lwe_in_ct = vec![0u64; lwe_in_n + 1];
375 encrypt_lwe(
376 &mut lwe_in_ct,
377 &encoded_m,
378 lwe_sk_in.values(),
379 &modq_op,
380 &mut rng,
381 );
382
383 // key switch from lwe_sk_in to lwe_sk_out
384 let decomposer = DefaultDecomposer::new(1u64 << logq, logb, d_ks);
385 let mut lwe_out_ct = vec![0u64; lwe_out_n + 1];
386 lwe_key_switch(
387 &mut lwe_out_ct,
388 &lwe_in_ct,
389 &ksk.data,
390 &modq_op,
391 &decomposer,
392 );
393
394 // decrypt lwe_out_ct using lwe_sk_out
395 let encoded_m_back = decrypt_lwe(&lwe_out_ct, &lwe_sk_out.values(), &modq_op);
396 let m_back = ((((encoded_m_back as f64) * ((1 << logp) as f64)) / q as f64).round()
397 as u64)
398 % (1u64 << logp);
399 let noise =
400 measure_noise_lwe(&lwe_out_ct, lwe_sk_out.values(), &modq_op, &encoded_m);
401 println!("Noise: {noise}");
402 assert_eq!(m, m_back, "Expected {m} but got {m_back}");
403 // dbg!(m, m_back);
404 // dbg!(encoded_m, encoded_m_back);
405 }
406 }
407 }
408}