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