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::tglwe::TGLWE;
 10use crate::tlwe::{PublicKey, SecretKey, TLWE};
 11
 12#[derive(Clone, Debug)]
 13pub struct TLev<const K: usize>(pub(crate) Vec<TLWE<K>>);
 14
 15impl<const K: usize> TLev<K> {
 16    pub fn encode<const T: u64>(m: &Rq<T, 1>) -> T64 {
 17        let coeffs = m.coeffs();
 18        T64(coeffs[0].0) // N=1, so take the only coeff
 19    }
 20    pub fn decode<const T: u64>(p: &T64) -> Rq<T, 1> {
 21        Rq::<T, 1>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
 22    }
 23    pub fn encrypt(
 24        mut rng: impl Rng,
 25        beta: u32,
 26        l: u32,
 27        pk: &PublicKey<K>,
 28        m: &T64,
 29    ) -> Result<Self> {
 30        let tlev: Vec<TLWE<K>> = (1..l + 1)
 31            .map(|i| {
 32                TLWE::<K>::encrypt(&mut rng, pk, &(*m * (u64::MAX / beta.pow(i as u32) as u64)))
 33            })
 34            .collect::<Result<Vec<_>>>()?;
 35
 36        Ok(Self(tlev))
 37    }
 38    pub fn encrypt_s(
 39        mut rng: impl Rng,
 40        _beta: u32, // TODO rm, and make beta=2 always
 41        l: u32,
 42        sk: &SecretKey<K>,
 43        m: &T64,
 44    ) -> Result<Self> {
 45        let tlev: Vec<TLWE<K>> = (1..l as u64 + 1)
 46            .map(|i| {
 47                let aux = if i < 64 {
 48                    *m * (u64::MAX / (1u64 << i))
 49                } else {
 50                    // 1<<64 would overflow, and anyways we're dividing u64::MAX
 51                    // by it, which would be equal to 1
 52                    *m
 53                };
 54                TLWE::<K>::encrypt_s(&mut rng, sk, &aux)
 55            })
 56            .collect::<Result<Vec<_>>>()?;
 57
 58        Ok(Self(tlev))
 59    }
 60
 61    pub fn decrypt(&self, sk: &SecretKey<K>, beta: u32) -> T64 {
 62        let pt = self.0[0].decrypt(sk);
 63        pt.mul_div_round(beta as u64, u64::MAX)
 64    }
 65}
 66// TODO review u64::MAX, since is -1 of the value we actually want
 67
 68impl<const K: usize> TLev<K> {
 69    pub fn iter(&self) -> std::slice::Iter<TLWE<K>> {
 70        self.0.iter()
 71    }
 72}
 73
 74// dot product between a TLev and Vec<T64>, usually Vec<T64> comes from a
 75// decomposition of T64
 76// TLev * Vec<T64> --> TLWE
 77impl<const K: usize> Mul<Vec<T64>> for TLev<K> {
 78    type Output = TLWE<K>;
 79    fn mul(self, v: Vec<T64>) -> Self::Output {
 80        assert_eq!(self.0.len(), v.len());
 81
 82        // l TLWES
 83        let tlwes: Vec<TLWE<K>> = self.0;
 84        let r: TLWE<K> = zip_eq(v, tlwes).map(|(a_d_i, glwe_i)| glwe_i * a_d_i).sum();
 85        r
 86    }
 87}
 88
 89#[cfg(test)]
 90mod tests {
 91    use anyhow::Result;
 92    use rand::distributions::Uniform;
 93
 94    use super::*;
 95
 96    #[test]
 97    fn test_encrypt_decrypt() -> Result<()> {
 98        const T: u64 = 2; // plaintext modulus
 99        const K: usize = 16;
100        type S = TLev<K>;
101
102        let beta: u32 = 2;
103        let l: u32 = 16;
104
105        let mut rng = rand::thread_rng();
106        let msg_dist = Uniform::new(0_u64, T);
107
108        for _ in 0..200 {
109            let (sk, pk) = TLWE::<K>::new_key(&mut rng)?;
110
111            let m: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
112            let p: T64 = S::encode::<T>(&m); // plaintext
113
114            let c = S::encrypt(&mut rng, beta, l, &pk, &p)?;
115            let p_recovered = c.decrypt(&sk, beta);
116            let m_recovered = S::decode::<T>(&p_recovered);
117
118            assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
119        }
120
121        Ok(())
122    }
123
124    #[test]
125    fn test_tlev_vect64_product() -> Result<()> {
126        const T: u64 = 2; // plaintext modulus
127        const K: usize = 16;
128
129        let beta: u32 = 2;
130        let l: u32 = 16;
131
132        let mut rng = rand::thread_rng();
133        let msg_dist = Uniform::new(0_u64, T);
134
135        for _ in 0..200 {
136            let (sk, pk) = TLWE::<K>::new_key(&mut rng)?;
137
138            let m1: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
139            let m2: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
140            let p1: T64 = TLev::<K>::encode::<T>(&m1);
141            let p2: T64 = TLev::<K>::encode::<T>(&m2);
142
143            let c1 = TLev::<K>::encrypt(&mut rng, beta, l, &pk, &p1)?;
144            let c2 = p2.decompose(beta, l);
145
146            let c3 = c1 * c2;
147
148            let p_recovered = c3.decrypt(&sk);
149            let m_recovered = TLev::<K>::decode::<T>(&p_recovered);
150
151            assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m_recovered);
152        }
153
154        Ok(())
155    }
156}