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