1use anyhow::Result;
 2use rand::Rng;
 3use rand_distr::{Normal, Uniform};
 4use std::ops::{Add, Mul};
 5
 6use arith::{Ring, Rq, TR};
 7
 8use crate::glwe::{PublicKey, SecretKey, GLWE};
 9
10const ERR_SIGMA: f64 = 3.2;
11
12// l GLWEs
13#[derive(Clone, Debug)]
14pub struct GLev<const Q: u64, const N: usize, const K: usize>(pub(crate) Vec<GLWE<Q, N, K>>);
15
16impl<const Q: u64, const N: usize, const K: usize> GLev<Q, N, K> {
17    pub fn encode<const T: u64>(m: &Rq<T, N>) -> Rq<Q, N> {
18        m.remodule::<Q>()
19    }
20    pub fn decode<const T: u64>(p: &Rq<Q, N>) -> Rq<T, N> {
21        p.remodule::<T>()
22    }
23    pub fn encrypt(
24        mut rng: impl Rng,
25        beta: u32,
26        l: u32,
27        pk: &PublicKey<Q, N, K>,
28        m: &Rq<Q, N>,
29    ) -> Result<Self> {
30        let glev: Vec<GLWE<Q, N, K>> = (1..l + 1)
31            .map(|i| {
32                GLWE::<Q, N, K>::encrypt(&mut rng, pk, &(*m * (Q / beta.pow(i as u32) as u64)))
33            })
34            .collect::<Result<Vec<_>>>()?;
35
36        Ok(Self(glev))
37    }
38    pub fn encrypt_s(
39        mut rng: impl Rng,
40        beta: u32,
41        l: u32,
42        sk: &SecretKey<Q, N, K>,
43        m: &Rq<Q, N>,
44    ) -> Result<Self> {
45        let glev: Vec<GLWE<Q, N, K>> = (1..l + 1)
46            .map(|i| {
47                GLWE::<Q, N, K>::encrypt_s(&mut rng, sk, &(*m * (Q / beta.pow(i as u32) as u64)))
48            })
49            .collect::<Result<Vec<_>>>()?;
50
51        Ok(Self(glev))
52    }
53
54    pub fn decrypt<const T: u64>(&self, sk: &SecretKey<Q, N, K>, beta: u32) -> Rq<Q, N> {
55        let pt = self.0[0].decrypt(sk);
56        pt.mul_div_round(beta as u64, Q)
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use anyhow::Result;
63    use rand::distributions::Uniform;
64
65    use super::*;
66
67    #[test]
68    fn test_encrypt_decrypt() -> Result<()> {
69        const Q: u64 = 2u64.pow(16) + 1;
70        const N: usize = 128;
71        const T: u64 = 2; // plaintext modulus
72        const K: usize = 16;
73        type S = GLev<Q, N, K>;
74
75        let beta: u32 = 2;
76        let l: u32 = 16;
77
78        let mut rng = rand::thread_rng();
79
80        for _ in 0..200 {
81            let (sk, pk) = GLWE::<Q, N, K>::new_key(&mut rng)?;
82
83            let msg_dist = Uniform::new(0_u64, T);
84            let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
85            let p: Rq<Q, N> = S::encode::<T>(&m); // plaintext
86
87            let c = S::encrypt(&mut rng, beta, l, &pk, &p)?;
88            let p_recovered = c.decrypt::<T>(&sk, beta);
89            let m_recovered = S::decode::<T>(&p_recovered);
90
91            assert_eq!(m, m_recovered);
92        }
93
94        Ok(())
95    }
96}