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::tglwe::{PublicKey, SecretKey, TGLWE};
 10use gfhe::glwe::{Param, GLWE};
 11
 12/// vector of length K+1 = ([K * TGLev], [1 * TGLev])
 13#[derive(Clone, Debug)]
 14pub struct TGGSW(pub(crate) Vec<TGLev>, TGLev);
 15
 16impl TGGSW {
 17    pub fn encrypt_s(
 18        mut rng: impl Rng,
 19        param: &Param,
 20        beta: u32,
 21        l: u32,
 22        sk: &SecretKey,
 23        m: &Tn,
 24    ) -> Result<Self> {
 25        debug_assert_eq!(sk.0 .0.k, param.k);
 26
 27        let a: Vec<TGLev> = (0..param.k)
 28            .map(|i| TGLev::encrypt_s(&mut rng, param, beta, l, sk, &(&-sk.0 .0.r[i].clone() * m)))
 29            // TODO rm clone
 30            .collect::<Result<Vec<_>>>()?;
 31        let b: TGLev = TGLev::encrypt_s(&mut rng, &param, beta, l, sk, m)?;
 32        Ok(Self(a, b))
 33    }
 34
 35    pub fn decrypt(&self, sk: &SecretKey, beta: u32) -> Tn {
 36        self.1.decrypt(sk, beta)
 37    }
 38
 39    pub fn cmux(bit: Self, ct1: TGLWE, ct2: TGLWE) -> TGLWE {
 40        ct1.clone() + (bit * (ct2 - ct1))
 41    }
 42}
 43
 44/// External product tggsw x tglwe
 45impl Mul<TGLWE> for TGGSW {
 46    type Output = TGLWE;
 47
 48    fn mul(self, tglwe: TGLWE) -> TGLWE {
 49        let beta: u32 = 2;
 50        let l: u32 = 64; // TODO wip
 51
 52        let tglwe_ab: Vec<Tn> = [tglwe.0 .0.r.clone(), vec![tglwe.0 .1]].concat();
 53
 54        let tgsw_ab: Vec<TGLev> = [self.0.clone(), vec![self.1]].concat();
 55        assert_eq!(tgsw_ab.len(), tglwe_ab.len());
 56
 57        let r: TGLWE = zip_eq(tgsw_ab, tglwe_ab)
 58            .map(|(tlev_i, tglwe_i)| tlev_i * tglwe_i.decompose(beta, l))
 59            .sum();
 60        r
 61    }
 62}
 63
 64#[derive(Clone, Debug)]
 65pub struct TGLev(pub(crate) Vec<TGLWE>);
 66
 67impl TGLev {
 68    pub fn encode(param: &Param, m: &Rq) -> Tn {
 69        debug_assert_eq!(param.t, m.param.q); // plaintext modulus
 70
 71        Tn {
 72            param: param.ring,
 73            coeffs: m.coeffs().iter().map(|c_i| T64(c_i.v)).collect(),
 74        }
 75    }
 76    pub fn decode(param: &Param, p: &Tn) -> Rq {
 77        Rq::from_vec_u64(&param.pt(), p.coeffs().iter().map(|c| c.0).collect())
 78    }
 79    pub fn encrypt(
 80        mut rng: impl Rng,
 81        param: &Param,
 82        beta: u32,
 83        l: u32,
 84        pk: &PublicKey,
 85        m: &Tn,
 86    ) -> Result<Self> {
 87        let tlev: Vec<TGLWE> = (1..l + 1)
 88            .map(|i| {
 89                TGLWE::encrypt(
 90                    &mut rng,
 91                    &param,
 92                    pk,
 93                    &(m * &(u64::MAX / beta.pow(i as u32) as u64)),
 94                )
 95            })
 96            .collect::<Result<Vec<_>>>()?;
 97
 98        Ok(Self(tlev))
 99    }
100    pub fn encrypt_s(
101        mut rng: impl Rng,
102        param: &Param,
103        _beta: u32, // TODO rm, and make beta=2 always
104        l: u32,
105        sk: &SecretKey,
106        m: &Tn,
107    ) -> Result<Self> {
108        let tlev: Vec<TGLWE> = (1..l as u64 + 1)
109            .map(|i| {
110                let aux = if i < 64 {
111                    m * &(u64::MAX / (1u64 << i))
112                } else {
113                    // 1<<64 would overflow, and anyways we're dividing u64::MAX
114                    // by it, which would be equal to 1
115                    m.clone() // TODO rm clone
116                };
117                TGLWE::encrypt_s(&mut rng, &param, sk, &aux)
118            })
119            .collect::<Result<Vec<_>>>()?;
120
121        Ok(Self(tlev))
122    }
123
124    pub fn decrypt(&self, sk: &SecretKey, beta: u32) -> Tn {
125        let pt = self.0[0].decrypt(sk);
126        pt.mul_div_round(beta as u64, u64::MAX)
127    }
128}
129
130impl TGLev {
131    pub fn iter(&self) -> std::slice::Iter<TGLWE> {
132        self.0.iter()
133    }
134}
135
136// dot product between a TGLev and Vec<Tn<N>>, usually Vec<Tn<N>> comes from a
137// decomposition of Tn<N>
138// TGLev * Vec<Tn<N>> --> TGLWE
139impl Mul<Vec<Tn>> for TGLev {
140    type Output = TGLWE;
141    fn mul(self, v: Vec<Tn>) -> Self::Output {
142        assert_eq!(self.0.len(), v.len());
143
144        // l TGLWES
145        let tlwes: Vec<TGLWE> = self.0;
146        let r: TGLWE = zip_eq(v, tlwes).map(|(a_d_i, glwe_i)| glwe_i * a_d_i).sum();
147        r
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use anyhow::Result;
154    use rand::distr::Uniform;
155
156    use super::*;
157    #[test]
158    fn test_external_product() -> Result<()> {
159        let param = Param {
160            err_sigma: crate::ERR_SIGMA,
161            ring: RingParam { q: u64::MAX, n: 64 },
162            k: 4,
163            t: 16, // plaintext modulus
164        };
165
166        let beta: u32 = 2;
167        let l: u32 = 64;
168
169        let mut rng = rand::rng();
170        let msg_dist = Uniform::new(0_u64, param.t)?;
171
172        for _ in 0..50 {
173            let (sk, _) = TGLWE::new_key(&mut rng, &param)?;
174
175            let m1: Rq = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
176            let p1: Tn = TGLev::encode(&param, &m1);
177
178            let m2: Rq = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
179            let p2: Tn = TGLWE::encode(&param, &m2); // scaled by delta
180
181            let tgsw = TGGSW::encrypt_s(&mut rng, &param, beta, l, &sk, &p1)?;
182            let tlwe = TGLWE::encrypt_s(&mut rng, &param, &sk, &p2)?;
183
184            let res: TGLWE = tgsw * tlwe;
185
186            // let p_recovered = res.decrypt(&sk, beta);
187            let p_recovered = res.decrypt(&sk);
188            // downscaled by delta^-1
189            let res_recovered = TGLWE::decode(&param, &p_recovered);
190
191            // assert_eq!(m1 * m2, m_recovered);
192            assert_eq!((m1.to_r() * m2.to_r()).to_rq(param.t), res_recovered);
193        }
194
195        Ok(())
196    }
197}