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