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    pub fn cmux(bit: Self, ct1: TLWE<K>, ct2: TLWE<K>) -> TLWE<K> {
 40        ct1.clone() + (bit * (ct2 - ct1))
 41    }
 42}
 43
 44/// External product TGSW x TLWE
 45impl<const K: usize> Mul<TLWE<K>> for TGSW<K> {
 46    type Output = TLWE<K>;
 47
 48    fn mul(self, tlwe: TLWE<K>) -> TLWE<K> {
 49        let beta: u32 = 2;
 50        let l: u32 = 64; // TODO wip
 51
 52        // since N=1, each tlwe element is a vector of length=1, decomposed into
 53        // l elements, and we have K of them
 54        let tlwe_ab: Vec<T64> = [tlwe.0 .0 .0.clone(), vec![tlwe.0 .1]].concat();
 55
 56        let tgsw_ab: Vec<TLev<K>> = [self.0.clone(), vec![self.1]].concat();
 57        assert_eq!(tgsw_ab.len(), tlwe_ab.len());
 58
 59        let r: TLWE<K> = zip_eq(tgsw_ab, tlwe_ab)
 60            .map(|(tlev_i, tlwe_i)| tlev_i * tlwe_i.decompose(beta, l))
 61            .sum();
 62        r
 63    }
 64}
 65
 66#[cfg(test)]
 67mod tests {
 68    use anyhow::Result;
 69    use rand::distributions::Uniform;
 70
 71    use super::*;
 72
 73    #[test]
 74    fn test_encrypt_decrypt() -> Result<()> {
 75        const T: u64 = 2; // plaintext modulus
 76        const K: usize = 16;
 77        type S = TGSW<K>;
 78
 79        let beta: u32 = 2;
 80        let l: u32 = 16;
 81
 82        let mut rng = rand::thread_rng();
 83        let msg_dist = Uniform::new(0_u64, T);
 84
 85        for _ in 0..50 {
 86            let (sk, _) = TLWE::<K>::new_key(&mut rng)?;
 87
 88            let m: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
 89            let p: T64 = TLev::<K>::encode::<T>(&m); // plaintext
 90
 91            let c = S::encrypt_s(&mut rng, beta, l, &sk, &p)?;
 92            let p_recovered = c.decrypt(&sk, beta);
 93            let m_recovered = TLev::<K>::decode::<T>(&p_recovered);
 94
 95            assert_eq!(m, m_recovered);
 96        }
 97
 98        Ok(())
 99    }
100
101    #[test]
102    fn test_external_product() -> Result<()> {
103        const T: u64 = 2; // plaintext modulus
104        const K: usize = 32;
105
106        let beta: u32 = 2;
107        let l: u32 = 64;
108
109        let mut rng = rand::thread_rng();
110        let msg_dist = Uniform::new(0_u64, T);
111
112        for _ in 0..50 {
113            let (sk, _) = TLWE::<K>::new_key(&mut rng)?;
114
115            let m1: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
116            let p1: T64 = TLev::<K>::encode::<T>(&m1);
117
118            let m2: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
119            let p2: T64 = TLWE::<K>::encode::<T>(&m2); // scaled by delta
120
121            let tgsw = TGSW::<K>::encrypt_s(&mut rng, beta, l, &sk, &p1)?;
122            let tlwe = TLWE::<K>::encrypt_s(&mut rng, &sk, &p2)?;
123
124            let res: TLWE<K> = tgsw * tlwe;
125
126            // let p_recovered = res.decrypt(&sk, beta);
127            let p_recovered = res.decrypt(&sk);
128            // downscaled by delta^-1
129            let res_recovered = TLWE::<K>::decode::<T>(&p_recovered);
130
131            // assert_eq!(m1 * m2, m_recovered);
132            assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), res_recovered);
133        }
134
135        Ok(())
136    }
137
138    #[test]
139    fn test_cmux() -> Result<()> {
140        const T: u64 = 2; // plaintext modulus
141        const K: usize = 32;
142
143        let beta: u32 = 2;
144        let l: u32 = 64;
145
146        let mut rng = rand::thread_rng();
147        let msg_dist = Uniform::new(0_u64, T);
148
149        for _ in 0..50 {
150            let (sk, _) = TLWE::<K>::new_key(&mut rng)?;
151
152            let m1: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
153            let p1: T64 = TLWE::<K>::encode::<T>(&m1); // scaled by delta
154
155            let m2: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?;
156            let p2: T64 = TLWE::<K>::encode::<T>(&m2); // scaled by delta
157
158            for bit_raw in 0..2 {
159                let bit = TGSW::<K>::encrypt_s(&mut rng, beta, l, &sk, &T64(bit_raw))?;
160
161                let c1 = TLWE::<K>::encrypt_s(&mut rng, &sk, &p1)?;
162                let c2 = TLWE::<K>::encrypt_s(&mut rng, &sk, &p2)?;
163
164                let res: TLWE<K> = TGSW::cmux(bit, c1, c2);
165
166                let p_recovered = res.decrypt(&sk);
167                // downscaled by delta^-1
168                let res_recovered = TLWE::<K>::decode::<T>(&p_recovered);
169
170                if bit_raw == 0 {
171                    assert_eq!(m1, res_recovered);
172                } else {
173                    assert_eq!(m2, res_recovered);
174                }
175            }
176        }
177
178        Ok(())
179    }
180}