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, T64, TR};
 11use gfhe::{glwe, GLWE};
 12
 13use crate::tlev::TLev;
 14
 15pub type SecretKey<const N: usize, const K: usize> = glwe::SecretKey<Tn<N>, K>;
 16pub type PublicKey<const N: usize, const K: usize> = glwe::PublicKey<Tn<N>, K>;
 17
 18#[derive(Clone, Debug)]
 19pub struct TGLWE<const N: usize, const K: usize>(pub GLWE<Tn<N>, K>);
 20
 21impl<const N: usize, const K: usize> TGLWE<N, K> {
 22    pub fn zero() -> Self {
 23        Self(GLWE::<Tn<N>, K>::zero())
 24    }
 25
 26    pub fn new_key(rng: impl Rng) -> Result<(SecretKey<N, K>, PublicKey<N, K>)> {
 27        let (sk, pk) = GLWE::new_key(rng)?;
 28        // Ok((SecretKey(sk), PublicKey(pk)))
 29        Ok((sk, pk))
 30    }
 31
 32    pub fn encode<const P: u64>(m: &Rq<P, N>) -> Tn<N> {
 33        let delta = u64::MAX / P; // floored
 34        let coeffs = m.coeffs();
 35        Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
 36    }
 37    pub fn decode<const P: u64>(p: &Tn<N>) -> Rq<P, N> {
 38        let p = p.mul_div_round(P, u64::MAX);
 39        Rq::<P, N>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
 40    }
 41
 42    // encrypts with the given SecretKey (instead of PublicKey)
 43    pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<N, K>, p: &Tn<N>) -> Result<Self> {
 44        let glwe = GLWE::encrypt_s(rng, &sk, p)?;
 45        Ok(Self(glwe))
 46    }
 47    pub fn encrypt(rng: impl Rng, pk: &PublicKey<N, K>, p: &Tn<N>) -> Result<Self> {
 48        let glwe = GLWE::encrypt(rng, &pk, p)?;
 49        Ok(Self(glwe))
 50    }
 51    pub fn decrypt(&self, sk: &SecretKey<N, K>) -> Tn<N> {
 52        self.0.decrypt(&sk)
 53    }
 54}
 55
 56impl<const N: usize, const K: usize> Add<TGLWE<N, K>> for TGLWE<N, K> {
 57    type Output = Self;
 58    fn add(self, other: Self) -> Self {
 59        Self(self.0 + other.0)
 60    }
 61}
 62impl<const N: usize, const K: usize> AddAssign for TGLWE<N, K> {
 63    fn add_assign(&mut self, rhs: Self) {
 64        self.0 += rhs.0
 65    }
 66}
 67impl<const N: usize, const K: usize> Sum<TGLWE<N, K>> for TGLWE<N, K> {
 68    fn sum<I>(iter: I) -> Self
 69    where
 70        I: Iterator<Item = Self>,
 71    {
 72        let mut acc = TGLWE::<N, K>::zero();
 73        for e in iter {
 74            acc += e;
 75        }
 76        acc
 77    }
 78}
 79
 80impl<const N: usize, const K: usize> Sub<TGLWE<N, K>> for TGLWE<N, K> {
 81    type Output = Self;
 82    fn sub(self, other: Self) -> Self {
 83        Self(self.0 - other.0)
 84    }
 85}
 86
 87// plaintext addition
 88impl<const N: usize, const K: usize> Add<Tn<N>> for TGLWE<N, K> {
 89    type Output = Self;
 90    fn add(self, plaintext: Tn<N>) -> Self {
 91        let a: TR<Tn<N>, K> = self.0 .0;
 92        let b: Tn<N> = self.0 .1 + plaintext;
 93        Self(GLWE(a, b))
 94    }
 95}
 96// plaintext substraction
 97impl<const N: usize, const K: usize> Sub<Tn<N>> for TGLWE<N, K> {
 98    type Output = Self;
 99    fn sub(self, plaintext: Tn<N>) -> Self {
100        let a: TR<Tn<N>, K> = self.0 .0;
101        let b: Tn<N> = self.0 .1 - plaintext;
102        Self(GLWE(a, b))
103    }
104}
105// plaintext multiplication
106impl<const N: usize, const K: usize> Mul<Tn<N>> for TGLWE<N, K> {
107    type Output = Self;
108    fn mul(self, plaintext: Tn<N>) -> Self {
109        let a: TR<Tn<N>, K> = TR(self.0 .0 .0.iter().map(|r_i| *r_i * plaintext).collect());
110        let b: Tn<N> = self.0 .1 * plaintext;
111        Self(GLWE(a, b))
112    }
113}
114
115#[cfg(test)]
116mod tests {
117    use anyhow::Result;
118    use rand::distributions::Uniform;
119
120    use super::*;
121
122    #[test]
123    fn test_encrypt_decrypt() -> Result<()> {
124        const T: u64 = 128; // msg space (msg modulus)
125        const N: usize = 64;
126        const K: usize = 16;
127        type S = TGLWE<N, K>;
128
129        let mut rng = rand::thread_rng();
130        let msg_dist = Uniform::new(0_u64, T);
131
132        for _ in 0..200 {
133            let (sk, pk) = S::new_key(&mut rng)?;
134
135            let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
136            let p: Tn<N> = S::encode::<T>(&m);
137
138            let c = S::encrypt(&mut rng, &pk, &p)?;
139            let p_recovered = c.decrypt(&sk);
140            let m_recovered = S::decode::<T>(&p_recovered);
141
142            assert_eq!(m, m_recovered);
143
144            // same but using encrypt_s (with sk instead of pk))
145            let c = S::encrypt_s(&mut rng, &sk, &p)?;
146            let p_recovered = c.decrypt(&sk);
147            let m_recovered = S::decode::<T>(&p_recovered);
148
149            assert_eq!(m, m_recovered);
150        }
151
152        Ok(())
153    }
154
155    #[test]
156    fn test_addition() -> Result<()> {
157        const T: u64 = 128;
158        const N: usize = 64;
159        const K: usize = 16;
160        type S = TGLWE<N, K>;
161
162        let mut rng = rand::thread_rng();
163        let msg_dist = Uniform::new(0_u64, T);
164
165        for _ in 0..200 {
166            let (sk, pk) = S::new_key(&mut rng)?;
167
168            let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
169            let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
170            let p1: Tn<N> = S::encode::<T>(&m1); // plaintext
171            let p2: Tn<N> = S::encode::<T>(&m2); // plaintext
172
173            let c1 = S::encrypt(&mut rng, &pk, &p1)?;
174            let c2 = S::encrypt(&mut rng, &pk, &p2)?;
175
176            let c3 = c1 + c2;
177
178            let p3_recovered = c3.decrypt(&sk);
179            let m3_recovered = S::decode::<T>(&p3_recovered);
180
181            assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>());
182        }
183
184        Ok(())
185    }
186
187    #[test]
188    fn test_add_plaintext() -> Result<()> {
189        const T: u64 = 128;
190        const N: usize = 64;
191        const K: usize = 16;
192        type S = TGLWE<N, K>;
193
194        let mut rng = rand::thread_rng();
195        let msg_dist = Uniform::new(0_u64, T);
196
197        for _ in 0..200 {
198            let (sk, pk) = S::new_key(&mut rng)?;
199
200            let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
201            let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
202            let p1: Tn<N> = S::encode::<T>(&m1); // plaintext
203            let p2: Tn<N> = S::encode::<T>(&m2); // plaintext
204
205            let c1 = S::encrypt(&mut rng, &pk, &p1)?;
206
207            let c3 = c1 + p2;
208
209            let p3_recovered = c3.decrypt(&sk);
210            let m3_recovered = S::decode::<T>(&p3_recovered);
211
212            assert_eq!(m1 + m2, m3_recovered);
213        }
214
215        Ok(())
216    }
217
218    #[test]
219    fn test_mul_plaintext() -> Result<()> {
220        const T: u64 = 128;
221        const N: usize = 64;
222        const K: usize = 16;
223        type S = TGLWE<N, K>;
224
225        let mut rng = rand::thread_rng();
226        let msg_dist = Uniform::new(0_u64, T);
227
228        for _ in 0..200 {
229            let (sk, pk) = S::new_key(&mut rng)?;
230
231            let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
232            let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
233            let p1: Tn<N> = S::encode::<T>(&m1);
234            // don't scale up p2, set it directly from m2
235            let p2: Tn<N> = Tn(array::from_fn(|i| T64(m2.coeffs()[i].0)));
236
237            let c1 = S::encrypt(&mut rng, &pk, &p1)?;
238
239            let c3 = c1 * p2;
240
241            let p3_recovered: Tn<N> = c3.decrypt(&sk);
242            let m3_recovered = S::decode::<T>(&p3_recovered);
243            assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered);
244        }
245
246        Ok(())
247    }
248}