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