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 encrypt(
18        mut rng: impl Rng,
19        beta: u32,
20        l: u32,
21        pk: &PublicKey<Q, N, K>,
22        m: &Rq<Q, N>,
23        // delta: u64,
24    ) -> Result<Self> {
25        let glev: Vec<GLWE<Q, N, K>> = (0..l)
26            .map(|i| {
27                GLWE::<Q, N, K>::encrypt(&mut rng, pk, &(*m * (Q / beta.pow(i as u32) as u64)), 1)
28            })
29            .collect::<Result<Vec<_>>>()?;
30
31        Ok(Self(glev))
32    }
33    pub fn encrypt_s(
34        mut rng: impl Rng,
35        beta: u32,
36        l: u32,
37        sk: &SecretKey<Q, N, K>,
38        m: &Rq<Q, N>,
39        // delta: u64,
40    ) -> Result<Self> {
41        let glev: Vec<GLWE<Q, N, K>> = (0..l)
42            .map(|i| {
43                GLWE::<Q, N, K>::encrypt_s(&mut rng, sk, &(*m * (Q / beta.pow(i as u32) as u64)), 1)
44            })
45            .collect::<Result<Vec<_>>>()?;
46
47        Ok(Self(glev))
48    }
49
50    pub fn decrypt<const T: u64>(&self, sk: &SecretKey<Q, N, K>, delta: u64) -> Rq<Q, N> {
51        self.0[1].decrypt::<T>(sk, delta)
52    }
53}
54
55#[cfg(test)]
56mod tests {
57    use anyhow::Result;
58    use rand::distributions::Uniform;
59
60    use super::*;
61
62    #[test]
63    fn test_encrypt_decrypt() -> Result<()> {
64        const Q: u64 = 2u64.pow(16) + 1;
65        const N: usize = 128;
66        const T: u64 = 2; // plaintext modulus
67        const K: usize = 16;
68        type S = GLev<Q, N, K>;
69
70        let beta: u32 = 2;
71        let l: u32 = 16;
72
73        let delta: u64 = Q / T; // floored
74        let mut rng = rand::thread_rng();
75
76        for _ in 0..200 {
77            let (sk, pk) = GLWE::<Q, N, K>::new_key(&mut rng)?;
78
79            let msg_dist = Uniform::new(0_u64, T);
80            let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
81            let m: Rq<Q, N> = m.remodule::<Q>();
82
83            let c = S::encrypt(&mut rng, beta, l, &pk, &m)?;
84            let m_recovered = c.decrypt::<T>(&sk, delta);
85
86            assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
87        }
88
89        Ok(())
90    }
91}