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// external product TGSW x TLWE
 41impl<const K: usize> Mul<TLWE<K>> for TGSW<K> {
 42    type Output = TLWE<K>;
 43
 44    fn mul(self, tlwe: TLWE<K>) -> TLWE<K> {
 45        let beta: u32 = 2;
 46        let l: u32 = 64; // TODO wip
 47
 48        // since N=1, each tlwe element is a vector of length=1, decomposed into
 49        // l elements, and we have K of them
 50        let tlwe_ab: Vec<T64> = [tlwe.0 .0 .0.clone(), vec![tlwe.0 .1]].concat();
 51
 52        let tgsw_ab: Vec<TLev<K>> = [self.0.clone(), vec![self.1]].concat();
 53        assert_eq!(tgsw_ab.len(), tlwe_ab.len());
 54
 55        let r: TLWE<K> = zip_eq(tgsw_ab, tlwe_ab)
 56            .map(|(tlev_i, tlwe_i)| tlev_i * tlwe_i.decompose(beta, l))
 57            .sum();
 58        r
 59    }
 60}
 61
 62#[cfg(test)]
 63mod tests {
 64    use anyhow::Result;
 65    use rand::distributions::Uniform;
 66
 67    use super::*;
 68
 69    #[test]
 70    fn test_encrypt_decrypt() -> Result<()> {
 71        const T: u64 = 2; // plaintext modulus
 72        const K: usize = 16;
 73        type S = TGSW<K>;
 74
 75        let beta: u32 = 2;
 76        let l: u32 = 16;
 77
 78        let mut rng = rand::thread_rng();
 79        let msg_dist = Uniform::new(0_u64, T);
 80
 81        for _ in 0..50 {
 82            let (sk, _) = TLWE::<K>::new_key(&mut rng)?;
 83
 84            let m: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
 85            let p: T64 = TLev::<K>::encode::<T>(&m); // plaintext
 86
 87            let c = S::encrypt_s(&mut rng, beta, l, &sk, &p)?;
 88            let p_recovered = c.decrypt(&sk, beta);
 89            let m_recovered = TLev::<K>::decode::<T>(&p_recovered);
 90
 91            assert_eq!(m, m_recovered);
 92        }
 93
 94        Ok(())
 95    }
 96
 97    #[test]
 98    fn test_external_product() -> Result<()> {
 99        const T: u64 = 2; // plaintext modulus
100        const K: usize = 32;
101
102        let beta: u32 = 2;
103        let l: u32 = 64;
104
105        let mut rng = rand::thread_rng();
106        let msg_dist = Uniform::new(0_u64, T);
107
108        for i in 0..50 {
109            dbg!(&i);
110            let (sk, _) = TLWE::<K>::new_key(&mut rng)?;
111
112            let m1: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
113            let p1: T64 = TLev::<K>::encode::<T>(&m1);
114
115            let m2: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
116            let p2: T64 = TLWE::<K>::encode::<T>(&m2); // scaled by delta
117
118            let tgsw = TGSW::<K>::encrypt_s(&mut rng, beta, l, &sk, &p1)?;
119            let tlwe = TLWE::<K>::encrypt_s(&mut rng, &sk, &p2)?;
120
121            let res: TLWE<K> = tgsw * tlwe;
122
123            // let p_recovered = res.decrypt(&sk, beta);
124            let p_recovered = res.decrypt(&sk);
125            // downscaled by delta^-1
126            let res_recovered = TLWE::<K>::decode::<T>(&p_recovered);
127
128            // assert_eq!(m1 * m2, m_recovered);
129            assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), res_recovered);
130        }
131
132        Ok(())
133    }
134}