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