(TFHE): add TLWE encryption & decryption
@@ -3,7 +3,8 @@ members = [
"arith", "gfhe", "bfv", - "ckks" + "ckks", + "tfhe" ] resolver = "2"
@@ -1,9 +1,10 @@
# fhe-study Implementations from scratch done while studying some FHE papers; do not use in production. -- `arith`: contains $\mathbb{Z}_q$, $R_q=\mathbb{Z}_q[X]/(X^N+1)$ and $R=\mathbb{Z}[X]/(X^N+1)$ arithmetic implementations, together with the NTT implementation. +- `arith`: contains $\mathbb{Z}_q$, $R_q=\mathbb{Z}_q[X]/(X^N+1)$, $R=\mathbb{Z}[X]/(X^N+1)$, $\mathbb{T}_{Q}[X]/(X^N +1)$ arithmetic implementations, together with the NTT implementation. - `gfhe`: (gfhe=generalized-fhe) contains the structs and logic for RLWE, GLWE, GLev, GGSW, RGSW cryptosystems, and modulus switching and key switching methods, which can be used by concrete FHE schemes. - `bfv`: https://eprint.iacr.org/2012/144.pdf scheme implementation - `ckks`: https://eprint.iacr.org/2016/421.pdf scheme implementation +- `tfhe`: https://eprint.iacr.org/2018/421.pdf scheme implementation `cargo test --release`
@@ -30,4 +30,10 @@
fn from_vec(coeffs: Vec<Self::C>) -> Self; fn decompose(&self, beta: u32, l: u32) -> Vec<Self>; + + /// returns [ [(num/den) * self].round() ] mod q + /// ie. performs the multiplication and division over f64, and then it + /// rounds the result, only applying the mod Q (if the ring is mod Q) at the + /// end. + fn mul_div_round(&self, num: u64, den: u64) -> Self; }
@@ -43,6 +43,19 @@ fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
unimplemented!(); // array::from_fn(|i| self.coeffs[i].decompose(beta, l)) } + + // performs the multiplication and division over f64, and then it rounds the + // result, only applying the mod Q at the end + fn mul_div_round(&self, num: u64, den: u64) -> Self { + unimplemented!() + // fn mul_div_round<const Q: u64>(&self, num: u64, den: u64) -> crate::Rq<Q, N> { + // let r: Vec<f64> = self + // .coeffs() + // .iter() + // .map(|e| ((num as f64 * *e as f64) / den as f64).round()) + // .collect(); + // crate::Rq::<Q, N>::from_vec_f64(r) + } } impl<const Q: u64, const N: usize> From<crate::ring_nq::Rq<Q, N>> for R<N> {@@ -73,16 +86,6 @@ Self(coeffs)
} pub fn mul_by_i64(&self, s: i64) -> Self { Self(array::from_fn(|i| self.0[i] * s)) - } - // performs the multiplication and division over f64, and then it rounds the - // result, only applying the mod Q at the end - pub fn mul_div_round<const Q: u64>(&self, num: u64, den: u64) -> crate::Rq<Q, N> { - let r: Vec<f64> = self - .coeffs() - .iter() - .map(|e| ((num as f64 * *e as f64) / den as f64).round()) - .collect(); - crate::Rq::<Q, N>::from_vec_f64(r) } pub fn infinity_norm(&self) -> u64 {
@@ -70,6 +70,18 @@ .collect();
// convert it to Rq<Q,N> r.iter().map(|a_i| Self::from_vec(a_i.clone())).collect() } + + // returns [ [(num/den) * self].round() ] mod q + // ie. performs the multiplication and division over f64, and then it rounds the + // result, only applying the mod Q at the end + fn mul_div_round(&self, num: u64, den: u64) -> Self { + let r: Vec<f64> = self + .coeffs() + .iter() + .map(|e| ((num as f64 * e.0 as f64) / den as f64).round()) + .collect(); + Rq::<Q, N>::from_vec_f64(r) + } } impl<const Q: u64, const N: usize> From<crate::ring_n::R<N>> for Rq<Q, N> {@@ -228,17 +240,6 @@ let r: Vec<f64> = self
.coeffs() .iter() .map(|e| (e.0 as f64 / s as f64).round()) - .collect(); - Rq::<Q, N>::from_vec_f64(r) - } - // returns [ [(num/den) * self].round() ] mod q - // ie. performs the multiplication and division over f64, and then it rounds the - // result, only applying the mod Q at the end - pub fn mul_div_round(&self, num: u64, den: u64) -> Self { - let r: Vec<f64> = self - .coeffs() - .iter() - .map(|e| ((num as f64 * e.0 as f64) / den as f64).round()) .collect(); Rq::<Q, N>::from_vec_f64(r) }
@@ -49,6 +49,18 @@ .collect();
// convert it to Tn<N> r.iter().map(|a_i| Self::from_vec(a_i.clone())).collect() } + + /// returns [ [(num/den) * self].round() ] mod q + /// ie. performs the multiplication and division over f64, and then it rounds the + /// result, only applying the mod Q at the end + fn mul_div_round(&self, num: u64, den: u64) -> Self { + let r: Vec<T64> = self + .coeffs() + .iter() + .map(|e| T64(((num as f64 * e.0 as f64) / den as f64).round() as u64)) + .collect(); + Self::from_vec(r) + } } // apply mod (X^N+1)
@@ -1,12 +1,15 @@
//! This file implements the struct for an Tuple of Ring Rq elements and its -//! operations. +//! operations, which are performed element-wise. use anyhow::Result; use itertools::zip_eq; use rand::{distributions::Distribution, Rng}; use rand_distr::{Normal, Uniform}; use std::iter::Sum; -use std::{array, ops}; +use std::{ + array, + ops::{Add, Mul, Sub}, +}; use crate::Ring;@@ -34,7 +37,7 @@ // self.0.iter().map(|r| r.decompose(beta, l)).collect() // this is Vec<Vec<Vec<R::C>>>
} } -impl<R: Ring, const K: usize> ops::Add<TR<R, K>> for TR<R, K> { +impl<R: Ring, const K: usize> Add<TR<R, K>> for TR<R, K> { type Output = Self; fn add(self, other: Self) -> Self { Self(@@ -45,7 +48,7 @@ )
} } -impl<R: Ring, const K: usize> ops::Sub<TR<R, K>> for TR<R, K> { +impl<R: Ring, const K: usize> Sub<TR<R, K>> for TR<R, K> { type Output = Self; fn sub(self, other: Self) -> Self { Self(zip_eq(self.0, other.0).map(|(s, o)| s - o).collect())@@ -54,13 +57,13 @@ }
/// for (TR,TR), the Mul operation is defined as: /// for A, B \in R^k, result = Σ A_i * B_i \in R -impl<R: Ring, const K: usize> ops::Mul<TR<R, K>> for TR<R, K> { +impl<R: Ring, const K: usize> Mul<TR<R, K>> for TR<R, K> { type Output = R; fn mul(self, other: Self) -> R { zip_eq(self.0, other.0).map(|(s, o)| s * o).sum() } } -impl<R: Ring, const K: usize> ops::Mul<&TR<R, K>> for &TR<R, K> { +impl<R: Ring, const K: usize> Mul<&TR<R, K>> for &TR<R, K> { type Output = R; fn mul(self, other: &TR<R, K>) -> R { zip_eq(self.0.clone(), other.0.clone())@@ -71,13 +74,13 @@ }
/// for (TR, R), the Mul operation is defined as each element of TR is /// multiplied by R -impl<R: Ring, const K: usize> ops::Mul<R> for TR<R, K> { +impl<R: Ring, const K: usize> Mul<R> for TR<R, K> { type Output = TR<R, K>; fn mul(self, other: R) -> TR<R, K> { Self(self.0.iter().map(|s| s.clone() * other.clone()).collect()) } } -impl<R: Ring, const K: usize> ops::Mul<&R> for &TR<R, K> { +impl<R: Ring, const K: usize> Mul<&R> for &TR<R, K> { type Output = TR<R, K>; fn mul(self, other: &R) -> TR<R, K> { TR::<R, K>(self.0.iter().map(|s| s.clone() * other.clone()).collect())
@@ -14,17 +14,22 @@ #[derive(Clone, Debug)]
pub struct GLev<const Q: u64, const N: usize, const K: usize>(pub(crate) Vec<GLWE<Q, N, K>>); impl<const Q: u64, const N: usize, const K: usize> GLev<Q, N, K> { + pub fn encode<const T: u64>(m: &Rq<T, N>) -> Rq<Q, N> { + m.remodule::<Q>() + } + pub fn decode<const T: u64>(p: &Rq<Q, N>) -> Rq<T, N> { + p.remodule::<T>() + } pub fn encrypt( mut rng: impl Rng, beta: u32, l: u32, pk: &PublicKey<Q, N, K>, m: &Rq<Q, N>, - // delta: u64, ) -> Result<Self> { - let glev: Vec<GLWE<Q, N, K>> = (0..l) + let glev: Vec<GLWE<Q, N, K>> = (1..l + 1) .map(|i| { - GLWE::<Q, N, K>::encrypt(&mut rng, pk, &(*m * (Q / beta.pow(i as u32) as u64)), 1) + GLWE::<Q, N, K>::encrypt(&mut rng, pk, &(*m * (Q / beta.pow(i as u32) as u64))) }) .collect::<Result<Vec<_>>>()?;@@ -36,19 +41,19 @@ beta: u32,
l: u32, sk: &SecretKey<Q, N, K>, m: &Rq<Q, N>, - // delta: u64, ) -> Result<Self> { - let glev: Vec<GLWE<Q, N, K>> = (0..l) + let glev: Vec<GLWE<Q, N, K>> = (1..l + 1) .map(|i| { - GLWE::<Q, N, K>::encrypt_s(&mut rng, sk, &(*m * (Q / beta.pow(i as u32) as u64)), 1) + GLWE::<Q, N, K>::encrypt_s(&mut rng, sk, &(*m * (Q / beta.pow(i as u32) as u64))) }) .collect::<Result<Vec<_>>>()?; Ok(Self(glev)) } - pub fn decrypt<const T: u64>(&self, sk: &SecretKey<Q, N, K>, delta: u64) -> Rq<Q, N> { - self.0[1].decrypt::<T>(sk, delta) + pub fn decrypt<const T: u64>(&self, sk: &SecretKey<Q, N, K>, beta: u32) -> Rq<Q, N> { + let pt = self.0[0].decrypt(sk); + pt.mul_div_round(beta as u64, Q) } }@@ -70,7 +75,6 @@
let beta: u32 = 2; let l: u32 = 16; - let delta: u64 = Q / T; // floored let mut rng = rand::thread_rng(); for _ in 0..200 {@@ -78,12 +82,13 @@ let (sk, pk) = GLWE::<Q, N, K>::new_key(&mut rng)?;
let msg_dist = Uniform::new(0_u64, T); let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; - let m: Rq<Q, N> = m.remodule::<Q>(); + let p: Rq<Q, N> = S::encode::<T>(&m); // plaintext - let c = S::encrypt(&mut rng, beta, l, &pk, &m)?; - let m_recovered = c.decrypt::<T>(&sk, delta); + let c = S::encrypt(&mut rng, beta, l, &pk, &p)?; + let p_recovered = c.decrypt::<T>(&sk, beta); + let m_recovered = S::decode::<T>(&p_recovered); - assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>()); + assert_eq!(m, m_recovered); } Ok(())
@@ -80,29 +80,30 @@ .sum();
r } + // scale up + pub fn encode<const T: u64>(m: &Rq<T, N>) -> Rq<Q, N> { + let m = m.remodule::<Q>(); + let delta = Q / T; // floored + m * delta + } + // scale down + pub fn decode<const T: u64>(p: &Rq<Q, N>) -> Rq<T, N> { + let r = p.mul_div_round(T, Q); + r.remodule::<T>() + } + // encrypts with the given SecretKey (instead of PublicKey) - pub fn encrypt_s( - mut rng: impl Rng, - sk: &SecretKey<Q, N, K>, - m: &Rq<Q, N>, - // TODO delta not as input - delta: u64, - ) -> Result<Self> { + pub fn encrypt_s(mut rng: impl Rng, sk: &SecretKey<Q, N, K>, m: &Rq<Q, N>) -> Result<Self> { let Xi_key = Uniform::new(0_f64, 2_f64); let Xi_err = Normal::new(0_f64, ERR_SIGMA)?; let a: TR<Rq<Q, N>, K> = TR::rand(&mut rng, Xi_key); let e = Rq::<Q, N>::rand(&mut rng, Xi_err); - let b: Rq<Q, N> = (&a * &sk.0) + *m * delta + e; + let b: Rq<Q, N> = (&a * &sk.0) + *m + e; Ok(Self(a, b)) } - pub fn encrypt( - mut rng: impl Rng, - pk: &PublicKey<Q, N, K>, - m: &Rq<Q, N>, - delta: u64, - ) -> Result<Self> { + pub fn encrypt(mut rng: impl Rng, pk: &PublicKey<Q, N, K>, m: &Rq<Q, N>) -> Result<Self> { let Xi_key = Uniform::new(0_f64, 2_f64); let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;@@ -111,15 +112,14 @@
let e0 = Rq::<Q, N>::rand(&mut rng, Xi_err); let e1 = TR::<Rq<Q, N>, K>::rand(&mut rng, Xi_err); - let b: Rq<Q, N> = pk.0 * u + *m * delta + e0; + let b: Rq<Q, N> = pk.0 * u + *m + e0; let d: TR<Rq<Q, N>, K> = &pk.1 * &u + e1; Ok(Self(d, b)) } - pub fn decrypt<const T: u64>(&self, sk: &SecretKey<Q, N, K>, delta: u64) -> Rq<Q, N> { + pub fn decrypt(&self, sk: &SecretKey<Q, N, K>) -> Rq<Q, N> { let (d, b): (TR<Rq<Q, N>, K>, Rq<Q, N>) = (self.0.clone(), self.1); let r: Rq<Q, N> = b - &d * &sk.0; - let r = r.mul_div_round(T, Q); r }@@ -215,7 +215,6 @@ const T: u64 = 32; // plaintext modulus
const K: usize = 16; type S = GLWE<Q, N, K>; - let delta: u64 = Q / T; // floored let mut rng = rand::thread_rng(); for _ in 0..200 {@@ -223,16 +222,18 @@ let (sk, pk) = S::new_key(&mut rng)?;
let msg_dist = Uniform::new(0_u64, T); let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; - let m: Rq<Q, N> = m.remodule::<Q>(); + let p = S::encode::<T>(&m); // plaintext - let c = S::encrypt(&mut rng, &pk, &m, delta)?; - let m_recovered = c.decrypt::<T>(&sk, delta); + let c = S::encrypt(&mut rng, &pk, &p)?; + let p_recovered = c.decrypt(&sk); + let m_recovered = S::decode::<T>(&p_recovered); assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>()); // same but using encrypt_s (with sk instead of pk)) - let c = S::encrypt_s(&mut rng, &sk, &m, delta)?; - let m_recovered = c.decrypt::<T>(&sk, delta); + let c = S::encrypt_s(&mut rng, &sk, &p)?; + let p_recovered = c.decrypt(&sk); + let m_recovered = S::decode::<T>(&p_recovered); assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>()); }@@ -248,7 +249,6 @@ const T: u64 = 20;
const K: usize = 16; type S = GLWE<Q, N, K>; - let delta: u64 = Q / T; // floored let mut rng = rand::thread_rng(); for _ in 0..200 {@@ -257,15 +257,16 @@
let msg_dist = Uniform::new(0_u64, T); let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; - let m1: Rq<Q, N> = m1.remodule::<Q>(); - let m2: Rq<Q, N> = m2.remodule::<Q>(); + let p1: Rq<Q, N> = S::encode::<T>(&m1); // plaintext + let p2: Rq<Q, N> = S::encode::<T>(&m2); // plaintext - let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?; - let c2 = S::encrypt(&mut rng, &pk, &m2, delta)?; + let c1 = S::encrypt(&mut rng, &pk, &p1)?; + let c2 = S::encrypt(&mut rng, &pk, &p2)?; let c3 = c1 + c2; - let m3_recovered = c3.decrypt::<T>(&sk, delta); + let p3_recovered = c3.decrypt(&sk); + let m3_recovered = S::decode::<T>(&p3_recovered); assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>()); }@@ -281,7 +282,6 @@ const T: u64 = 32;
const K: usize = 16; type S = GLWE<Q, N, K>; - let delta: u64 = Q / T; // floored let mut rng = rand::thread_rng(); for _ in 0..200 {@@ -290,17 +290,17 @@
let msg_dist = Uniform::new(0_u64, T); let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; - let m1: Rq<Q, N> = m1.remodule::<Q>(); - let m2: Rq<Q, N> = m2.remodule::<Q>(); - let m2_scaled: Rq<Q, N> = m2 * delta; + let p1: Rq<Q, N> = S::encode::<T>(&m1); // plaintext + let p2: Rq<Q, N> = S::encode::<T>(&m2); // plaintext - let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?; + let c1 = S::encrypt(&mut rng, &pk, &p1)?; - let c3 = c1 + m2_scaled; + let c3 = c1 + p2; - let m3_recovered = c3.decrypt::<T>(&sk, delta); + let p3_recovered = c3.decrypt(&sk); + let m3_recovered = S::decode::<T>(&p3_recovered); - assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>()); + assert_eq!((m1 + m2).remodule::<T>(), m3_recovered); } Ok(())@@ -314,7 +314,6 @@ const T: u64 = 4;
const K: usize = 16; type S = GLWE<Q, N, K>; - let delta: u64 = Q / T; // floored let mut rng = rand::thread_rng(); for _ in 0..200 {@@ -323,14 +322,15 @@
let msg_dist = Uniform::new(0_u64, T); let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; - let m1: Rq<Q, N> = m1.remodule::<Q>(); - let m2: Rq<Q, N> = m2.remodule::<Q>(); - let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?; + let p1: Rq<Q, N> = S::encode::<T>(&m1); // plaintext + let p2: Rq<Q, N> = m2.remodule::<Q>(); - let c3 = c1 * m2; + let c1 = S::encrypt(&mut rng, &pk, &p1)?; + + let c3 = c1 * p2; - let m3_recovered: Rq<Q, N> = c3.decrypt::<T>(&sk, delta); - let m3_recovered: Rq<T, N> = m3_recovered.remodule::<T>(); + let p3_recovered: Rq<Q, N> = c3.decrypt(&sk); + let m3_recovered = S::decode::<T>(&p3_recovered); assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered); }@@ -360,17 +360,18 @@ let (sk, pk) = S::new_key(&mut rng)?;
let msg_dist = Uniform::new(0_u64, T); let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; - let m: Rq<Q, N> = m.remodule::<Q>(); + let p = S::encode::<T>(&m); // plaintext - let c = S::encrypt(&mut rng, &pk, &m, delta)?; + let c = S::encrypt(&mut rng, &pk, &p)?; // let c = S::encrypt_s(&mut rng, &sk, &m, delta)?; let c2 = c.mod_switch::<P>(); let sk2: SecretKey<P, N, K> = SecretKey(TR(sk.0 .0.iter().map(|s_i| s_i.remodule::<P>()).collect())); - let delta2: u64 = ((P as f64 * delta as f64) / Q as f64).round() as u64; + // let delta2: u64 = ((P as f64 * delta as f64) / Q as f64).round() as u64; - let m_recovered = c2.decrypt::<T>(&sk2, delta2); + let p_recovered = c2.decrypt(&sk2); + let m_recovered = GLWE::<P, N, K>::decode::<T>(&p_recovered); assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>()); }@@ -389,7 +390,6 @@
let beta: u32 = 2; let l: u32 = 16; - let delta: u64 = Q / T; // floored let mut rng = rand::thread_rng(); let (sk, pk) = S::new_key(&mut rng)?;@@ -399,21 +399,23 @@ let ksk = S::new_ksk(&mut rng, beta, l, &sk, &sk2)?;
let msg_dist = Uniform::new(0_u64, T); let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; - let m: Rq<Q, N> = m.remodule::<Q>(); + let p: Rq<Q, N> = S::encode::<T>(&m); // plaintext - let c = S::encrypt_s(&mut rng, &sk, &m, delta)?; + let c = S::encrypt_s(&mut rng, &sk, &p)?; let c2 = c.key_switch(beta, l, &ksk); // decrypt with the 2nd secret key - let m_recovered = c2.decrypt::<T>(&sk2, delta); - assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>()); + let p_recovered = c2.decrypt(&sk2); + let m_recovered = S::decode::<T>(&p_recovered); + assert_eq!(m, m_recovered); // do the same but now encrypting with pk - // let c = S::encrypt(&mut rng, &pk, &m, delta)?; + // let c = S::encrypt(&mut rng, &pk, &p)?; // let c2 = c.key_switch(beta, l, &ksk); - // let m_recovered = c2.decrypt::<T>(&sk2, delta); - // assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>()); + // let p_recovered = c2.decrypt(&sk2); + // let m_recovered = S::decode::<T>(&p_recovered); + // assert_eq!(m, m_recovered); Ok(()) }
@@ -0,0 +1,12 @@
+[package] +name = "tfhe" +version = "0.1.0" +edition = "2024" + +[dependencies] +anyhow = { workspace = true } +rand = { workspace = true } +rand_distr = { workspace = true } +itertools = { workspace = true } + +arith = { path="../arith" }
@@ -0,0 +1,8 @@
+//! Implementation of TFHE https://eprint.iacr.org/2018/421.pdf +#![allow(non_snake_case)] +#![allow(non_upper_case_globals)] +#![allow(non_camel_case_types)] +#![allow(clippy::upper_case_acronyms)] +#![allow(dead_code)] // TMP + +pub mod tlwe;
@@ -0,0 +1,120 @@
+use anyhow::Result; +use itertools::zip_eq; +use rand::distributions::Standard; +use rand::Rng; +use rand_distr::{Normal, Uniform}; +use std::array; +use std::iter::Sum; +use std::ops::{Add, AddAssign, Mul, Sub}; + +use arith::{Ring, Rq, Tn, Zq, T64, TR}; + +const ERR_SIGMA: f64 = 3.2; + +#[derive(Clone, Debug)] +pub struct TLWE<const K: usize>(TR<Tn<1>, K>, Tn<1>); + +#[derive(Clone, Debug)] +pub struct SecretKey<const K: usize>(TR<Tn<1>, K>); +#[derive(Clone, Debug)] +pub struct PublicKey<const K: usize>(Tn<1>, TR<Tn<1>, K>); + +impl<const K: usize> TLWE<K> { + pub fn zero() -> Self { + Self(TR::zero(), Tn::zero()) + } + + pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> { + let Xi_key = Uniform::new(0_f64, 2_f64); + let Xi_err = Normal::new(0_f64, ERR_SIGMA)?; + + let s: TR<Tn<1>, K> = TR::rand(&mut rng, Xi_key); + let a: TR<Tn<1>, K> = TR::rand(&mut rng, Standard); + let e = Tn::rand(&mut rng, Xi_err); + + let pk: PublicKey<K> = PublicKey((&a * &s) + e, a); + Ok((SecretKey(s), pk)) + } + + pub fn encode<const P: u64>(m: Rq<P, 1>) -> Tn<1> { + let delta = u64::MAX / P; // floored + let coeffs = m.coeffs(); + Tn(array::from_fn(|i| T64(coeffs[i].0 * delta))) + } + pub fn decode<const P: u64>(p: Tn<1>) -> Rq<P, 1> { + let p = p.mul_div_round(P, u64::MAX); + Rq::<P, 1>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect()) + } + + // encrypts with the given SecretKey (instead of PublicKey) + pub fn encrypt_s(mut rng: impl Rng, sk: &SecretKey<K>, m: &Tn<1>) -> Result<Self> { + let Xi_key = Uniform::new(0_f64, 2_f64); + let Xi_err = Normal::new(0_f64, ERR_SIGMA)?; + + let a: TR<Tn<1>, K> = TR::rand(&mut rng, Xi_key); + let e = Tn::rand(&mut rng, Xi_err); + + let b: Tn<1> = (&a * &sk.0) + *m + e; + Ok(Self(a, b)) + } + pub fn encrypt(mut rng: impl Rng, pk: &PublicKey<K>, m: &Tn<1>) -> Result<Self> { + let Xi_key = Uniform::new(0_f64, 2_f64); + let Xi_err = Normal::new(0_f64, ERR_SIGMA)?; + + let u: Tn<1> = Tn::rand(&mut rng, Xi_key); + + let e0: Tn<1> = Tn::rand(&mut rng, Xi_err); + let e1 = TR::<Tn<1>, K>::rand(&mut rng, Xi_err); + + let b: Tn<1> = pk.0 * u + *m + e0; + let d: TR<Tn<1>, K> = &pk.1 * &u + e1; + + Ok(Self(d, b)) + } + pub fn decrypt(&self, sk: &SecretKey<K>) -> Tn<1> { + let (d, b): (TR<Tn<1>, K>, Tn<1>) = (self.0.clone(), self.1); + b - &d * &sk.0 + } +} + +#[cfg(test)] +mod tests { + use anyhow::Result; + use rand::distributions::Uniform; + + use super::*; + + #[test] + fn test_encrypt_decrypt() -> Result<()> { + const T: u64 = 32; // plaintext modulus + const K: usize = 16; + type S = TLWE<K>; + + let mut rng = rand::thread_rng(); + + for _ in 0..200 { + let (sk, pk) = S::new_key(&mut rng)?; + + let msg_dist = Uniform::new(0_u64, T); + let m = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?; + dbg!(&m); + let p: Tn<1> = S::encode::<T>(m); + dbg!(&p); + + let c = S::encrypt(&mut rng, &pk, &p)?; + let p_recovered = c.decrypt(&sk); + let m_recovered = S::decode::<T>(p_recovered); + + assert_eq!(m, m_recovered); + + // same but using encrypt_s (with sk instead of pk)) + let c = S::encrypt_s(&mut rng, &sk, &p)?; + let p_recovered = c.decrypt(&sk); + let m_recovered = S::decode::<T>(p_recovered); + + assert_eq!(m, m_recovered); + } + + Ok(()) + } +}