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