1//! Implementation of TFHE https://eprint.iacr.org/2018/421.pdf
 2#![allow(non_snake_case)]
 3#![allow(non_upper_case_globals)]
 4#![allow(non_camel_case_types)]
 5#![allow(clippy::upper_case_acronyms)]
 6#![allow(dead_code)] // TMP
 7
 8use anyhow::Result;
 9use rand::Rng;
10use rand_distr::{Normal, Uniform};
11use std::array;
12
13use arith::{Ring, Rq, Tn, T64};
14use gfhe::{glwe, GLWE};
15
16pub mod tlev;
17pub mod tlwe;
18
19#[derive(Clone, Debug)]
20pub struct SecretKey<const K: usize>(glwe::SecretKey<Tn<1>, K>);
21#[derive(Clone, Debug)]
22pub struct PublicKey<const K: usize>(glwe::PublicKey<Tn<1>, K>);
23
24#[derive(Clone, Debug)]
25pub struct TLWE<const K: usize>(pub GLWE<Tn<1>, K>);
26
27impl<const K: usize> TLWE<K> {
28    pub fn new_key(rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> {
29        let (sk, pk) = GLWE::new_key(rng)?;
30        Ok((SecretKey(sk), PublicKey(pk)))
31    }
32
33    pub fn encode<const P: u64>(m: &Rq<P, 1>) -> Tn<1> {
34        let delta = u64::MAX / P; // floored
35        let coeffs = m.coeffs();
36        Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
37    }
38    pub fn decode<const P: u64>(p: &Tn<1>) -> Rq<P, 1> {
39        let p = p.mul_div_round(P, u64::MAX);
40        Rq::<P, 1>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
41    }
42    pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<K>, p: &Tn<1>) -> Result<Self> {
43        let glwe = GLWE::encrypt_s(rng, &sk.0, p)?;
44        Ok(Self(glwe))
45    }
46
47    pub fn encrypt(rng: impl Rng, pk: &PublicKey<K>, p: &Tn<1>) -> Result<Self> {
48        let glwe = GLWE::encrypt(rng, &pk.0, p)?;
49        Ok(Self(glwe))
50    }
51
52    pub fn decrypt(&self, sk: &SecretKey<K>) -> Tn<1> {
53        self.0.decrypt(&sk.0)
54    }
55}
56
57#[cfg(test)]
58mod tests {
59    use anyhow::Result;
60    use rand::distributions::Uniform;
61
62    use super::*;
63
64    #[test]
65    fn test_encrypt_decrypt() -> Result<()> {
66        const T: u64 = 128; // plaintext modulus
67        const K: usize = 16;
68        type S = TLWE<K>;
69
70        // let delta: u64 = Q / T; // floored
71        let mut rng = rand::thread_rng();
72
73        for _ in 0..200 {
74            let (sk, pk) = S::new_key(&mut rng)?;
75
76            let msg_dist = Uniform::new(0_f64, T as f64);
77            let m = Rq::<T, 1>::rand(&mut rng, msg_dist); // msg
78
79            // let m: Rq<Q, N> = m.remodule::<Q>();
80
81            let p = S::encode::<T>(&m); // plaintext
82            let c = S::encrypt(&mut rng, &pk, &p)?; // ciphertext
83            let p_recovered = c.decrypt(&sk);
84            let m_recovered = S::decode::<T>(&p_recovered);
85
86            assert_eq!(m, m_recovered);
87
88            // same but using encrypt_s (with sk instead of pk))
89            let c = S::encrypt_s(&mut rng, &sk, &p)?;
90            let p_recovered = c.decrypt(&sk);
91            let m_recovered = S::decode::<T>(&p_recovered);
92
93            assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
94        }
95
96        Ok(())
97    }
98}