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