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::tlev::TLev;
10use crate::tlwe::{PublicKey, SecretKey, TLWE};
11use gfhe::glwe::GLWE;
12
13/// vector of length K+1 = [K], [1]
14#[derive(Clone, Debug)]
15pub struct TGSW<const K: usize>(pub(crate) Vec<TLev<K>>, TLev<K>);
16
17impl<const K: usize> TGSW<K> {
18    pub fn encrypt_s(
19        mut rng: impl Rng,
20        beta: u32,
21        l: u32,
22        sk: &SecretKey<K>,
23        m: &T64,
24    ) -> Result<Self> {
25        let a: Vec<TLev<K>> = (0..K)
26            .map(|i| TLev::encrypt_s(&mut rng, beta, l, sk, &(-sk.0 .0[i] * *m)))
27            .collect::<Result<Vec<_>>>()?;
28        let b: TLev<K> = TLev::encrypt_s(&mut rng, beta, l, sk, m)?;
29        Ok(Self(a, b))
30    }
31
32    pub fn decrypt(&self, sk: &SecretKey<K>, beta: u32) -> T64 {
33        self.1.decrypt(sk, beta)
34    }
35    pub fn from_tlwe(_tlwe: TLWE<K>) -> Self {
36        todo!()
37    }
38}
39
40#[cfg(test)]
41mod tests {
42    use anyhow::Result;
43    use rand::distributions::Uniform;
44
45    use super::*;
46
47    #[test]
48    fn test_encrypt_decrypt() -> Result<()> {
49        const T: u64 = 2; // plaintext modulus
50        const K: usize = 16;
51        type S = TGSW<K>;
52
53        let beta: u32 = 2;
54        let l: u32 = 16;
55
56        let mut rng = rand::thread_rng();
57        let msg_dist = Uniform::new(0_u64, T);
58
59        for _ in 0..50 {
60            let (sk, _) = TLWE::<K>::new_key(&mut rng)?;
61
62            let m: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
63            let p: T64 = TLev::<K>::encode::<T>(&m); // plaintext
64
65            let c = S::encrypt_s(&mut rng, beta, l, &sk, &p)?;
66            let p_recovered = c.decrypt(&sk, beta);
67            let m_recovered = TLev::<K>::decode::<T>(&p_recovered);
68
69            assert_eq!(m, m_recovered);
70        }
71
72        Ok(())
73    }
74}