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