1use anyhow::Result;
  2use itertools::zip_eq;
  3use rand::Rng;
  4use std::ops::Mul;
  5
  6use arith::{Ring, RingParam, Rq, T64};
  7
  8use crate::tlwe::{PublicKey, SecretKey, TLWE};
  9use gfhe::glwe::Param;
 10
 11#[derive(Clone, Debug)]
 12pub struct TLev(pub(crate) Vec<TLWE>);
 13
 14impl TLev {
 15    pub fn encode(param: &Param, m: &Rq) -> T64 {
 16        assert_eq!(m.param.n, 1);
 17        assert_eq!(param.t, m.param.q);
 18
 19        let coeffs = m.coeffs();
 20        T64(coeffs[0].v) // N=1, so take the only coeff
 21    }
 22    pub fn decode(param: &Param, p: &T64) -> Rq {
 23        Rq::from_vec_u64(
 24            &RingParam { q: param.t, n: 1 },
 25            p.coeffs().iter().map(|c| c.0).collect(),
 26        )
 27    }
 28    pub fn encrypt(
 29        mut rng: impl Rng,
 30        param: &Param,
 31        beta: u32,
 32        l: u32,
 33        pk: &PublicKey,
 34        m: &T64,
 35    ) -> Result<Self> {
 36        debug_assert_eq!(pk.1.k, param.k);
 37
 38        let tlev: Vec<TLWE> = (1..l as u64 + 1)
 39            .map(|i| {
 40                let aux = if i < 64 {
 41                    *m * (u64::MAX / (1u64 << i))
 42                } else {
 43                    // 1<<64 would overflow, and anyways we're dividing u64::MAX
 44                    // by it, which would be equal to 1
 45                    *m
 46                };
 47                TLWE::encrypt(&mut rng, param, pk, &aux)
 48            })
 49            .collect::<Result<Vec<_>>>()?;
 50
 51        Ok(Self(tlev))
 52    }
 53    pub fn encrypt_s(
 54        mut rng: impl Rng,
 55        param: &Param,
 56        _beta: u32, // TODO rm, and make beta=2 always
 57        l: u32,
 58        sk: &SecretKey,
 59        m: &T64,
 60    ) -> Result<Self> {
 61        debug_assert_eq!(sk.0 .0.k, param.k);
 62
 63        let tlev: Vec<TLWE> = (1..l as u64 + 1)
 64            .map(|i| {
 65                let aux = if i < 64 {
 66                    *m * (u64::MAX / (1u64 << i))
 67                } else {
 68                    // 1<<64 would overflow, and anyways we're dividing u64::MAX
 69                    // by it, which would be equal to 1
 70                    *m
 71                };
 72                TLWE::encrypt_s(&mut rng, &param, sk, &aux)
 73            })
 74            .collect::<Result<Vec<_>>>()?;
 75
 76        Ok(Self(tlev))
 77    }
 78
 79    pub fn decrypt(&self, sk: &SecretKey, beta: u32) -> T64 {
 80        let pt = self.0[0].decrypt(sk);
 81        pt.mul_div_round(beta as u64, u64::MAX)
 82    }
 83}
 84// TODO review u64::MAX, since is -1 of the value we actually want
 85
 86impl TLev {
 87    pub fn iter(&self) -> std::slice::Iter<TLWE> {
 88        self.0.iter()
 89    }
 90}
 91
 92// dot product between a TLev and Vec<T64>, usually Vec<T64> comes from a
 93// decomposition of T64
 94// TLev * Vec<T64> --> TLWE
 95impl Mul<Vec<T64>> for TLev {
 96    type Output = TLWE;
 97    fn mul(self, v: Vec<T64>) -> Self::Output {
 98        assert_eq!(self.0.len(), v.len());
 99
100        // l TLWES
101        let tlwes: Vec<TLWE> = self.0;
102        let r: TLWE = zip_eq(v, tlwes).map(|(a_d_i, glwe_i)| glwe_i * a_d_i).sum();
103        r
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use anyhow::Result;
110    use rand::distr::Uniform;
111
112    use super::*;
113
114    #[test]
115    fn test_encrypt_decrypt() -> Result<()> {
116        let param = Param {
117            err_sigma: crate::ERR_SIGMA,
118            ring: RingParam { q: u64::MAX, n: 1 },
119            k: 16,
120            t: 2, // plaintext modulus
121        };
122
123        let beta: u32 = 2;
124        let l: u32 = 16;
125
126        let mut rng = rand::rng();
127        let msg_dist = Uniform::new(0_u64, param.t)?;
128
129        for _ in 0..200 {
130            let (sk, pk) = TLWE::new_key(&mut rng, &param)?;
131
132            let m: Rq = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
133            let p: T64 = TLev::encode(&param, &m); // plaintext
134
135            let c = TLev::encrypt(&mut rng, &param, beta, l, &pk, &p)?;
136            let p_recovered = c.decrypt(&sk, beta);
137            let m_recovered = TLev::decode(&param, &p_recovered);
138
139            assert_eq!(m.remodule(param.t), m_recovered.remodule(param.t));
140        }
141
142        Ok(())
143    }
144
145    #[test]
146    fn test_tlev_vect64_product() -> Result<()> {
147        let param = Param {
148            err_sigma: 0.1, // WIP
149            ring: RingParam { q: u64::MAX, n: 1 },
150            k: 16,
151            t: 2, // plaintext modulus
152        };
153
154        let beta: u32 = 2;
155        // let l: u32 = 16;
156        let l: u32 = 64;
157
158        let mut rng = rand::rng();
159        let msg_dist = Uniform::new(0_u64, param.t)?;
160
161        for _ in 0..200 {
162            let (sk, pk) = TLWE::new_key(&mut rng, &param)?;
163
164            let m1: Rq = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
165            let m2: Rq = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
166            let p1: T64 = TLev::encode(&param, &m1);
167            let p2: T64 = TLev::encode(&param, &m2);
168
169            let c1 = TLev::encrypt(&mut rng, &param, beta, l, &pk, &p1)?;
170            let c2 = p2.decompose(beta, l);
171
172            let c3 = c1 * c2;
173
174            let p_recovered = c3.decrypt(&sk);
175            let m_recovered = TLev::decode(&param, &p_recovered);
176
177            assert_eq!((m1.to_r() * m2.to_r()).to_rq(param.t), m_recovered);
178        }
179
180        Ok(())
181    }
182}