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