fhe-study/gfhe/src/glwe.rs
gfhe/src/glwe.rsRaw
1//! Generalized LWE.
2//!
3
4use anyhow::Result;
5use itertools::zip_eq;
6use rand::Rng;
7use rand_distr::{Normal, Uniform};
8use std::iter::Sum;
9use std::ops::{Add, AddAssign, Mul, Sub};
10
11use arith::{Ring, Rq, Zq, TR};
12
13use crate::glev::GLev;
14
15const ERR_SIGMA: f64 = 3.2;
16
17/// GLWE implemented over the `Ring` trait, so that it can be also instantiated
18/// over the Torus polynomials 𝕋_<N,q>[X] = 𝕋_q[X]/ (X^N+1).
19#[derive(Clone, Debug)]
20pub struct GLWE<R: Ring, const K: usize>(TR<R, K>, R);
21
22#[derive(Clone, Debug)]
23pub struct SecretKey<R: Ring, const K: usize>(TR<R, K>);
24#[derive(Clone, Debug)]
25pub struct PublicKey<R: Ring, const K: usize>(R, TR<R, K>);
26
27// K GLevs, each KSK_i=l GLWEs
28#[derive(Clone, Debug)]
29pub struct KSK<R: Ring, const K: usize>(Vec<GLev<R, K>>);
30
31impl<R: Ring, const K: usize> GLWE<R, K> {
32 pub fn zero() -> Self {
33 Self(TR::zero(), R::zero())
34 }
35
36 pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<R, K>, PublicKey<R, K>)> {
37 let Xi_key = Uniform::new(0_f64, 2_f64);
38 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
39
40 let s: TR<R, K> = TR::rand(&mut rng, Xi_key);
41 let a: TR<R, K> = TR::rand(&mut rng, Uniform::new(0_f64, R::Q as f64));
42 let e = R::rand(&mut rng, Xi_err);
43
44 let pk: PublicKey<R, K> = PublicKey((&a * &s) + e, a);
45 Ok((SecretKey(s), pk))
46 }
47
48 pub fn new_ksk(
49 mut rng: impl Rng,
50 beta: u32,
51 l: u32,
52 sk: &SecretKey<R, K>,
53 new_sk: &SecretKey<R, K>,
54 ) -> Result<KSK<R, K>> {
55 let r: Vec<GLev<R, K>> = (0..K)
56 .into_iter()
57 .map(|i|
58 // treat sk_i as the msg being encrypted
59 GLev::<R, K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0 .0[i]))
60 .collect::<Result<Vec<_>>>()?;
61
62 Ok(KSK(r))
63 }
64 pub fn key_switch(&self, beta: u32, l: u32, ksk: &KSK<R, K>) -> Self {
65 let (a, b): (TR<R, K>, R) = (self.0.clone(), self.1);
66
67 let lhs: GLWE<R, K> = GLWE(TR::zero(), b);
68
69 // K iterations, ksk.0 contains K times GLev
70 let rhs: GLWE<R, K> = zip_eq(a.0, ksk.0.clone())
71 .map(|(a_i, ksk_i)| Self::dot_prod(a_i.decompose(beta, l), ksk_i))
72 .sum();
73
74 lhs - rhs
75 }
76 // note: a_decomp is of length N
77 fn dot_prod(a_decomp: Vec<R>, ksk_i: GLev<R, K>) -> GLWE<R, K> {
78 // l times GLWES
79 let glwes: Vec<GLWE<R, K>> = ksk_i.0;
80
81 // l iterations
82 let r: GLWE<R, K> = zip_eq(a_decomp, glwes)
83 .map(|(a_d_i, glwe_i)| glwe_i * a_d_i)
84 .sum();
85 r
86 }
87
88 // encrypts with the given SecretKey (instead of PublicKey)
89 pub fn encrypt_s(
90 mut rng: impl Rng,
91 sk: &SecretKey<R, K>,
92 m: &R, // already scaled
93 ) -> Result<Self> {
94 let Xi_key = Uniform::new(0_f64, 2_f64);
95 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
96
97 let a: TR<R, K> = TR::rand(&mut rng, Xi_key);
98 let e = R::rand(&mut rng, Xi_err);
99
100 let b: R = (&a * &sk.0) + *m + e;
101 Ok(Self(a, b))
102 }
103 pub fn encrypt(
104 mut rng: impl Rng,
105 pk: &PublicKey<R, K>,
106 m: &R, // already scaled
107 ) -> Result<Self> {
108 let Xi_key = Uniform::new(0_f64, 2_f64);
109 let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
110
111 let u: R = R::rand(&mut rng, Xi_key);
112
113 let e0 = R::rand(&mut rng, Xi_err);
114 let e1 = TR::<R, K>::rand(&mut rng, Xi_err);
115
116 let b: R = pk.0.clone() * u.clone() + *m + e0;
117 let d: TR<R, K> = &pk.1 * &u + e1;
118
119 Ok(Self(d, b))
120 }
121 // returns m' not downscaled
122 pub fn decrypt(&self, sk: &SecretKey<R, K>) -> R {
123 let (d, b): (TR<R, K>, R) = (self.0.clone(), self.1);
124 let p: R = b - &d * &sk.0;
125 p
126 }
127}
128
129// Methods for when Ring=Rq<Q,N>
130impl<const Q: u64, const N: usize, const K: usize> GLWE<Rq<Q, N>, K> {
131 // scale up
132 pub fn encode<const T: u64>(m: &Rq<T, N>) -> Rq<Q, N> {
133 let m = m.remodule::<Q>();
134 let delta = Q / T; // floored
135 m * delta
136 }
137 // scale down
138 pub fn decode<const T: u64>(m: &Rq<Q, N>) -> Rq<T, N> {
139 let r = m.mul_div_round(T, Q);
140 let r: Rq<T, N> = r.remodule::<T>();
141 r
142 }
143 pub fn mod_switch<const P: u64>(&self) -> GLWE<Rq<P, N>, K> {
144 let a: TR<Rq<P, N>, K> = TR(self
145 .0
146 .0
147 .iter()
148 .map(|r| r.mod_switch::<P>())
149 .collect::<Vec<_>>());
150 let b: Rq<P, N> = self.1.mod_switch::<P>();
151 GLWE(a, b)
152 }
153}
154
155impl<R: Ring, const K: usize> Add<GLWE<R, K>> for GLWE<R, K> {
156 type Output = Self;
157 fn add(self, other: Self) -> Self {
158 let a: TR<R, K> = self.0 + other.0;
159 let b: R = self.1 + other.1;
160 Self(a, b)
161 }
162}
163
164impl<R: Ring, const K: usize> Add<R> for GLWE<R, K> {
165 type Output = Self;
166 fn add(self, plaintext: R) -> Self {
167 let a: TR<R, K> = self.0;
168 let b: R = self.1 + plaintext;
169 Self(a, b)
170 }
171}
172impl<R: Ring, const K: usize> AddAssign for GLWE<R, K> {
173 fn add_assign(&mut self, rhs: Self) {
174 for i in 0..K {
175 self.0 .0[i] = self.0 .0[i].clone() + rhs.0 .0[i].clone();
176 }
177 self.1 = self.1.clone() + rhs.1.clone();
178 }
179}
180impl<R: Ring, const K: usize> Sum<GLWE<R, K>> for GLWE<R, K> {
181 fn sum<I>(iter: I) -> Self
182 where
183 I: Iterator<Item = Self>,
184 {
185 let mut acc = GLWE::<R, K>::zero();
186 for e in iter {
187 acc += e;
188 }
189 acc
190 }
191}
192
193impl<R: Ring, const K: usize> Sub<GLWE<R, K>> for GLWE<R, K> {
194 type Output = Self;
195 fn sub(self, other: Self) -> Self {
196 let a: TR<R, K> = self.0 - other.0;
197 let b: R = self.1 - other.1;
198 Self(a, b)
199 }
200}
201
202impl<R: Ring, const K: usize> Mul<R> for GLWE<R, K> {
203 type Output = Self;
204 fn mul(self, plaintext: R) -> Self {
205 let a: TR<R, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
206 let b: R = self.1 * plaintext;
207 Self(a, b)
208 }
209}
210// for when R = Rq<Q,N>
211// impl<const Q: u64, const N: usize, const K: usize> Mul<Rq<Q, N>> for GLWE<Rq<Q, N>, K> {
212// type Output = Self;
213// fn mul(self, plaintext: Rq<Q, N>) -> Self {
214// // first compute the NTT for plaintext, to avoid computing it at each
215// // iteration, speeding up the multiplications
216// let mut plaintext = plaintext.clone();
217// plaintext.compute_evals();
218//
219// let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
220// let b: Rq<Q, N> = self.1 * plaintext;
221// Self(a, b)
222// }
223// }
224
225// impl<R: Ring, const K: usize> Mul<R::C> for GLWE<R, K>
226// // where
227// // // R: std::ops::Mul<<R as arith::Ring>::C>,
228// // // Vec<R>: FromIterator<<R as Mul<<R as arith::Ring>::C>>::Output>,
229// // Vec<R>: FromIterator<<R as Mul<<R as arith::Ring>::C>>::Output>,
230// {
231// type Output = Self;
232// fn mul(self, e: R::C) -> Self {
233// let a: TR<R, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e.clone()).collect());
234// let b: R = self.1 * e.clone();
235// Self(a, b)
236// }
237// }
238
239// impl<const Q: u64, const N: usize, const K: usize> Mul<Zq<Q>> for GLWE<Q, N, K> {
240// type Output = Self;
241// fn mul(self, e: Zq<Q>) -> Self {
242// let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e).collect());
243// let b: Rq<Q, N> = self.1 * e;
244// Self(a, b)
245// }
246// }
247
248#[cfg(test)]
249mod tests {
250 use anyhow::Result;
251 use rand::distributions::Uniform;
252
253 use super::*;
254
255 #[test]
256 fn test_encrypt_decrypt() -> Result<()> {
257 const Q: u64 = 2u64.pow(16) + 1;
258 const N: usize = 128;
259 const T: u64 = 32; // plaintext modulus
260 const K: usize = 16;
261 type S = GLWE<Rq<Q, N>, K>;
262
263 let mut rng = rand::thread_rng();
264
265 for _ in 0..200 {
266 let (sk, pk) = S::new_key(&mut rng)?;
267
268 let msg_dist = Uniform::new(0_u64, T);
269 let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; // msg
270 // let m: Rq<Q, N> = m.remodule::<Q>();
271
272 let p = S::encode::<T>(&m); // plaintext
273 let c = S::encrypt(&mut rng, &pk, &p)?; // ciphertext
274 let p_recovered = c.decrypt(&sk);
275 let m_recovered = S::decode::<T>(&p_recovered);
276
277 assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
278
279 // same but using encrypt_s (with sk instead of pk))
280 let c = S::encrypt_s(&mut rng, &sk, &p)?;
281 let p_recovered = c.decrypt(&sk);
282 let m_recovered = S::decode::<T>(&p_recovered);
283
284 assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
285 }
286
287 Ok(())
288 }
289
290 use arith::{Tn, T64};
291 use std::array;
292 pub fn t_encode<const P: u64>(m: &Rq<P, 4>) -> Tn<4> {
293 let delta = u64::MAX / P; // floored
294 let coeffs = m.coeffs();
295 Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
296 }
297 pub fn t_decode<const P: u64>(p: &Tn<4>) -> Rq<P, 4> {
298 let p = p.mul_div_round(P, u64::MAX);
299 Rq::<P, 4>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
300 }
301 #[test]
302 fn test_encrypt_decrypt_torus() -> Result<()> {
303 const N: usize = 128;
304 const T: u64 = 32; // plaintext modulus
305 const K: usize = 16;
306 type S = GLWE<Tn<4>, K>;
307
308 let mut rng = rand::thread_rng();
309
310 for _ in 0..200 {
311 let (sk, pk) = S::new_key(&mut rng)?;
312
313 let msg_dist = Uniform::new(0_f64, T as f64);
314 let m = Rq::<T, 4>::rand(&mut rng, msg_dist); // msg
315
316 let p = t_encode::<T>(&m); // plaintext
317 let c = S::encrypt(&mut rng, &pk, &p)?; // ciphertext
318 let p_recovered = c.decrypt(&sk);
319 let m_recovered = t_decode::<T>(&p_recovered);
320
321 assert_eq!(m, m_recovered);
322
323 // same but using encrypt_s (with sk instead of pk))
324 let c = S::encrypt_s(&mut rng, &sk, &p)?;
325 let p_recovered = c.decrypt(&sk);
326 let m_recovered = t_decode::<T>(&p_recovered);
327
328 assert_eq!(m, m_recovered);
329 }
330
331 Ok(())
332 }
333
334 #[test]
335 fn test_addition() -> Result<()> {
336 const Q: u64 = 2u64.pow(16) + 1;
337 const N: usize = 128;
338 const T: u64 = 20;
339 const K: usize = 16;
340 type S = GLWE<Rq<Q, N>, K>;
341
342 let mut rng = rand::thread_rng();
343
344 for _ in 0..200 {
345 let (sk, pk) = S::new_key(&mut rng)?;
346
347 let msg_dist = Uniform::new(0_u64, T);
348 let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
349 let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
350 let p1: Rq<Q, N> = S::encode::<T>(&m1); // plaintext
351 let p2: Rq<Q, N> = S::encode::<T>(&m2); // plaintext
352
353 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
354 let c2 = S::encrypt(&mut rng, &pk, &p2)?;
355
356 let c3 = c1 + c2;
357
358 let p3_recovered = c3.decrypt(&sk);
359 let m3_recovered = S::decode::<T>(&p3_recovered);
360
361 assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>());
362 }
363
364 Ok(())
365 }
366
367 #[test]
368 fn test_add_plaintext() -> Result<()> {
369 const Q: u64 = 2u64.pow(16) + 1;
370 const N: usize = 128;
371 const T: u64 = 32;
372 const K: usize = 16;
373 type S = GLWE<Rq<Q, N>, K>;
374
375 let mut rng = rand::thread_rng();
376
377 for _ in 0..200 {
378 let (sk, pk) = S::new_key(&mut rng)?;
379
380 let msg_dist = Uniform::new(0_u64, T);
381 let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
382 let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
383 let p1: Rq<Q, N> = S::encode::<T>(&m1); // plaintext
384 let p2: Rq<Q, N> = S::encode::<T>(&m2); // plaintext
385
386 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
387
388 let c3 = c1 + p2;
389
390 let p3_recovered = c3.decrypt(&sk);
391 let m3_recovered = S::decode::<T>(&p3_recovered);
392
393 assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>());
394 }
395
396 Ok(())
397 }
398
399 #[test]
400 fn test_mul_plaintext() -> Result<()> {
401 const Q: u64 = 2u64.pow(16) + 1;
402 const N: usize = 16;
403 const T: u64 = 4;
404 const K: usize = 16;
405 type S = GLWE<Rq<Q, N>, K>;
406
407 let mut rng = rand::thread_rng();
408
409 for _ in 0..200 {
410 let (sk, pk) = S::new_key(&mut rng)?;
411
412 let msg_dist = Uniform::new(0_u64, T);
413 let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
414 let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
415 let p1: Rq<Q, N> = S::encode::<T>(&m1); // plaintext
416 let p2 = m2.remodule::<Q>(); // notice we don't encode (scale by delta)
417
418 let c1 = S::encrypt(&mut rng, &pk, &p1)?;
419
420 let c3 = c1 * p2;
421
422 let p3_recovered: Rq<Q, N> = c3.decrypt(&sk);
423 let m3_recovered: Rq<T, N> = S::decode::<T>(&p3_recovered);
424 assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered);
425 }
426
427 Ok(())
428 }
429
430 #[test]
431 fn test_mod_switch() -> Result<()> {
432 const Q: u64 = 2u64.pow(16) + 1;
433 const P: u64 = 2u64.pow(8) + 1;
434 // note: wip, Q and P chosen so that P/Q is an integer
435 const N: usize = 8;
436 const T: u64 = 4; // plaintext modulus, must be a prime or power of a prime
437 const K: usize = 16;
438 type S = GLWE<Rq<Q, N>, K>;
439
440 let mut rng = rand::thread_rng();
441
442 for _ in 0..200 {
443 let (sk, pk) = S::new_key(&mut rng)?;
444
445 let msg_dist = Uniform::new(0_u64, T);
446 let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
447
448 let p = S::encode::<T>(&m);
449 let c = S::encrypt(&mut rng, &pk, &p)?;
450
451 let c2: GLWE<Rq<P, N>, K> = c.mod_switch::<P>();
452 let sk2: SecretKey<Rq<P, N>, K> =
453 SecretKey(TR(sk.0 .0.iter().map(|s_i| s_i.remodule::<P>()).collect()));
454
455 let p_recovered = c2.decrypt(&sk2);
456 let m_recovered = GLWE::<Rq<P, N>, K>::decode::<T>(&p_recovered);
457
458 assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
459 }
460
461 Ok(())
462 }
463
464 #[test]
465 fn test_key_switch() -> Result<()> {
466 const Q: u64 = 2u64.pow(16) + 1;
467 const N: usize = 128;
468 const T: u64 = 2; // plaintext modulus
469 const K: usize = 16;
470 type S = GLWE<Rq<Q, N>, K>;
471
472 let beta: u32 = 2;
473 let l: u32 = 16;
474
475 let mut rng = rand::thread_rng();
476
477 let (sk, pk) = S::new_key(&mut rng)?;
478 let (sk2, _) = S::new_key(&mut rng)?;
479 // ksk to switch from sk to sk2
480 let ksk = S::new_ksk(&mut rng, beta, l, &sk, &sk2)?;
481
482 let msg_dist = Uniform::new(0_u64, T);
483 let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
484 let p = S::encode::<T>(&m); // plaintext
485 //
486 let c = S::encrypt_s(&mut rng, &sk, &p)?;
487
488 let c2 = c.key_switch(beta, l, &ksk);
489
490 // decrypt with the 2nd secret key
491 let p_recovered = c2.decrypt(&sk2);
492 let m_recovered = S::decode::<T>(&p_recovered);
493 assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
494
495 // do the same but now encrypting with pk
496 let c = S::encrypt(&mut rng, &pk, &p)?;
497 let c2 = c.key_switch(beta, l, &ksk);
498 let p_recovered = c2.decrypt(&sk2);
499 let m_recovered = S::decode::<T>(&p_recovered);
500 assert_eq!(m, m_recovered);
501
502 Ok(())
503 }
504}