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