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