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