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