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