1use anyhow::Result;
  2use itertools::zip_eq;
  3use rand::distributions::Standard;
  4use rand::Rng;
  5use rand_distr::{Normal, Uniform};
  6use std::array;
  7use std::iter::Sum;
  8use std::ops::{Add, AddAssign, Mul, Sub};
  9
 10use arith::{Ring, Rq, Tn, Zq, T64, TR};
 11
 12const ERR_SIGMA: f64 = 3.2;
 13
 14#[derive(Clone, Debug)]
 15pub struct TLWE<const K: usize>(TR<Tn<1>, K>, Tn<1>);
 16
 17#[derive(Clone, Debug)]
 18pub struct SecretKey<const K: usize>(TR<Tn<1>, K>);
 19#[derive(Clone, Debug)]
 20pub struct PublicKey<const K: usize>(Tn<1>, TR<Tn<1>, K>);
 21
 22impl<const K: usize> TLWE<K> {
 23    pub fn zero() -> Self {
 24        Self(TR::zero(), Tn::zero())
 25    }
 26
 27    pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> {
 28        let Xi_key = Uniform::new(0_f64, 2_f64);
 29        let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
 30
 31        let s: TR<Tn<1>, K> = TR::rand(&mut rng, Xi_key);
 32        let a: TR<Tn<1>, K> = TR::rand(&mut rng, Standard);
 33        let e = Tn::rand(&mut rng, Xi_err);
 34
 35        let pk: PublicKey<K> = PublicKey((&a * &s) + e, a);
 36        Ok((SecretKey(s), pk))
 37    }
 38
 39    pub fn encode<const P: u64>(m: Rq<P, 1>) -> Tn<1> {
 40        let delta = u64::MAX / P; // floored
 41        let coeffs = m.coeffs();
 42        Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
 43    }
 44    pub fn decode<const P: u64>(p: Tn<1>) -> Rq<P, 1> {
 45        let p = p.mul_div_round(P, u64::MAX);
 46        Rq::<P, 1>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
 47    }
 48
 49    // encrypts with the given SecretKey (instead of PublicKey)
 50    pub fn encrypt_s(mut rng: impl Rng, sk: &SecretKey<K>, m: &Tn<1>) -> Result<Self> {
 51        let Xi_key = Uniform::new(0_f64, 2_f64);
 52        let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
 53
 54        let a: TR<Tn<1>, K> = TR::rand(&mut rng, Xi_key);
 55        let e = Tn::rand(&mut rng, Xi_err);
 56
 57        let b: Tn<1> = (&a * &sk.0) + *m + e;
 58        Ok(Self(a, b))
 59    }
 60    pub fn encrypt(mut rng: impl Rng, pk: &PublicKey<K>, m: &Tn<1>) -> Result<Self> {
 61        let Xi_key = Uniform::new(0_f64, 2_f64);
 62        let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
 63
 64        let u: Tn<1> = Tn::rand(&mut rng, Xi_key);
 65
 66        let e0: Tn<1> = Tn::rand(&mut rng, Xi_err);
 67        let e1 = TR::<Tn<1>, K>::rand(&mut rng, Xi_err);
 68
 69        let b: Tn<1> = pk.0 * u + *m + e0;
 70        let d: TR<Tn<1>, K> = &pk.1 * &u + e1;
 71
 72        Ok(Self(d, b))
 73    }
 74    pub fn decrypt(&self, sk: &SecretKey<K>) -> Tn<1> {
 75        let (d, b): (TR<Tn<1>, K>, Tn<1>) = (self.0.clone(), self.1);
 76        b - &d * &sk.0
 77    }
 78}
 79
 80#[cfg(test)]
 81mod tests {
 82    use anyhow::Result;
 83    use rand::distributions::Uniform;
 84
 85    use super::*;
 86
 87    #[test]
 88    fn test_encrypt_decrypt() -> Result<()> {
 89        const T: u64 = 32; // plaintext modulus
 90        const K: usize = 16;
 91        type S = TLWE<K>;
 92
 93        let mut rng = rand::thread_rng();
 94
 95        for _ in 0..200 {
 96            let (sk, pk) = S::new_key(&mut rng)?;
 97
 98            let msg_dist = Uniform::new(0_u64, T);
 99            let m = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
100            dbg!(&m);
101            let p: Tn<1> = S::encode::<T>(m);
102            dbg!(&p);
103
104            let c = S::encrypt(&mut rng, &pk, &p)?;
105            let p_recovered = c.decrypt(&sk);
106            let m_recovered = S::decode::<T>(p_recovered);
107
108            assert_eq!(m, m_recovered);
109
110            // same but using encrypt_s (with sk instead of pk))
111            let c = S::encrypt_s(&mut rng, &sk, &p)?;
112            let p_recovered = c.decrypt(&sk);
113            let m_recovered = S::decode::<T>(p_recovered);
114
115            assert_eq!(m, m_recovered);
116        }
117
118        Ok(())
119    }
120}