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}