arnaucube
arnaucube/fhe-studyPublic Self-hosted repository
add TGSW encryption & decryption
arnaucube git@arnaucube.com
Sun, 27 Jul 2025 19:41:15 +0000
commit

f053e9a90422148a2035ad39ea4a0afa9a857ef4

parent

4f89caef1eab17c5ae90fc9be115f9ae1b150247

4 files changed, 136 insertions(+), 48 deletions(-)

jump to
M tfhe/src/lib.rs → tfhe/src/lib.rs

@@ -5,5 +5,6 @@ #![allow(non_camel_case_types)]

#![allow(clippy::upper_case_acronyms)] #![allow(dead_code)] // TMP +pub mod tgsw; pub mod tlev; pub mod tlwe;
A tfhe/src/tgsw.rs

@@ -0,0 +1,74 @@

+use anyhow::Result; +use itertools::zip_eq; +use rand::Rng; +use std::array; +use std::ops::{Add, Mul}; + +use arith::{Ring, Rq, Tn, T64, TR}; + +use crate::tlev::TLev; +use crate::tlwe::{PublicKey, SecretKey, TLWE}; +use gfhe::glwe::GLWE; + +/// vector of length K+1 = [K], [1] +#[derive(Clone, Debug)] +pub struct TGSW<const K: usize>(pub(crate) Vec<TLev<K>>, TLev<K>); + +impl<const K: usize> TGSW<K> { + pub fn encrypt_s( + mut rng: impl Rng, + beta: u32, + l: u32, + sk: &SecretKey<K>, + m: &T64, + ) -> Result<Self> { + let a: Vec<TLev<K>> = (0..K) + .map(|i| TLev::encrypt_s(&mut rng, beta, l, sk, &(-sk.0 .0[i] * *m))) + .collect::<Result<Vec<_>>>()?; + let b: TLev<K> = TLev::encrypt_s(&mut rng, beta, l, sk, m)?; + Ok(Self(a, b)) + } + + pub fn decrypt(&self, sk: &SecretKey<K>, beta: u32) -> T64 { + self.1.decrypt(sk, beta) + } + pub fn from_tlwe(_tlwe: TLWE<K>) -> Self { + todo!() + } +} + +#[cfg(test)] +mod tests { + use anyhow::Result; + use rand::distributions::Uniform; + + use super::*; + + #[test] + fn test_encrypt_decrypt() -> Result<()> { + const T: u64 = 2; // plaintext modulus + const K: usize = 16; + type S = TGSW<K>; + + let beta: u32 = 2; + let l: u32 = 16; + + let mut rng = rand::thread_rng(); + let msg_dist = Uniform::new(0_u64, T); + + for _ in 0..50 { + let (sk, _) = TLWE::<K>::new_key(&mut rng)?; + + let m: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?; + let p: T64 = TLev::<K>::encode::<T>(&m); // plaintext + + let c = S::encrypt_s(&mut rng, beta, l, &sk, &p)?; + let p_recovered = c.decrypt(&sk, beta); + let m_recovered = TLev::<K>::decode::<T>(&p_recovered); + + assert_eq!(m, m_recovered); + } + + Ok(()) + } +}
M tfhe/src/tlev.rs → tfhe/src/tlev.rs

@@ -1,4 +1,5 @@

use anyhow::Result; +use itertools::zip_eq; use rand::Rng; use std::array; use std::ops::{Add, Mul};

@@ -11,11 +12,11 @@ #[derive(Clone, Debug)]

pub struct TLev<const K: usize>(pub(crate) Vec<TLWE<K>>); impl<const K: usize> TLev<K> { - pub fn encode<const T: u64>(m: &Rq<T, 1>) -> Tn<1> { + pub fn encode<const T: u64>(m: &Rq<T, 1>) -> T64 { let coeffs = m.coeffs(); - Tn(array::from_fn(|i| T64(coeffs[i].0))) + T64(coeffs[0].0) // N=1, so take the only coeff } - pub fn decode<const T: u64>(p: &Tn<1>) -> Rq<T, 1> { + pub fn decode<const T: u64>(p: &T64) -> Rq<T, 1> { Rq::<T, 1>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect()) } pub fn encrypt(

@@ -23,7 +24,7 @@ mut rng: impl Rng,

beta: u32, l: u32, pk: &PublicKey<K>, - m: &Tn<1>, + m: &T64, ) -> Result<Self> { let tlev: Vec<TLWE<K>> = (1..l + 1) .map(|i| {

@@ -35,25 +36,33 @@ Ok(Self(tlev))

} pub fn encrypt_s( mut rng: impl Rng, - beta: u32, + _beta: u32, // TODO rm, and make beta=2 always l: u32, sk: &SecretKey<K>, - m: &Tn<1>, + m: &T64, ) -> Result<Self> { - let tlev: Vec<TLWE<K>> = (1..l + 1) + let tlev: Vec<TLWE<K>> = (1..l as u64 + 1) .map(|i| { - TLWE::<K>::encrypt_s(&mut rng, sk, &(*m * (u64::MAX / beta.pow(i as u32) as u64))) + let aux = if i < 64 { + *m * (u64::MAX / (1u64 << i)) + } else { + // 1<<64 would overflow, and anyways we're dividing u64::MAX + // by it, which would be equal to 1 + *m + }; + TLWE::<K>::encrypt_s(&mut rng, sk, &aux) }) .collect::<Result<Vec<_>>>()?; Ok(Self(tlev)) } - pub fn decrypt(&self, sk: &SecretKey<K>, beta: u32) -> Tn<1> { + pub fn decrypt(&self, sk: &SecretKey<K>, beta: u32) -> T64 { let pt = self.0[0].decrypt(sk); pt.mul_div_round(beta as u64, u64::MAX) } } +// TODO review u64::MAX, since is -1 of the value we actually want #[cfg(test)] mod tests {

@@ -78,7 +87,7 @@ for _ in 0..200 {

let (sk, pk) = TLWE::<K>::new_key(&mut rng)?; let m: Rq<T, 1> = Rq::rand_u64(&mut rng, msg_dist)?; - let p: Tn<1> = S::encode::<T>(&m); // plaintext + let p: T64 = S::encode::<T>(&m); // plaintext let c = S::encrypt(&mut rng, beta, l, &pk, &p)?; let p_recovered = c.decrypt(&sk, beta);
M tfhe/src/tlwe.rs → tfhe/src/tlwe.rs

@@ -10,47 +10,50 @@

use arith::{Ring, Rq, Tn, T64, TR}; use gfhe::{glwe, GLWE}; -const ERR_SIGMA: f64 = 3.2; +// #[derive(Clone, Debug)] +// pub struct SecretKey<const K: usize>(glwe::SecretKey<T64, K>); +pub type SecretKey<const K: usize> = glwe::SecretKey<T64, K>; -#[derive(Clone, Debug)] -pub struct SecretKey<const K: usize>(glwe::SecretKey<Tn<1>, K>); -#[derive(Clone, Debug)] -pub struct PublicKey<const K: usize>(glwe::PublicKey<Tn<1>, K>); +// #[derive(Clone, Debug)] +// pub struct PublicKey<const K: usize>(glwe::PublicKey<T64, K>); +pub type PublicKey<const K: usize> = glwe::PublicKey<T64, K>; #[derive(Clone, Debug)] -pub struct TLWE<const K: usize>(pub GLWE<Tn<1>, K>); +pub struct TLWE<const K: usize>(pub GLWE<T64, K>); impl<const K: usize> TLWE<K> { pub fn zero() -> Self { - Self(GLWE::<Tn<1>, K>::zero()) + Self(GLWE::<T64, K>::zero()) } pub fn new_key(rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> { let (sk, pk) = GLWE::new_key(rng)?; - Ok((SecretKey(sk), PublicKey(pk))) + // Ok((SecretKey(sk), PublicKey(pk))) + Ok((sk, pk)) } - pub fn encode<const P: u64>(m: &Rq<P, 1>) -> Tn<1> { + pub fn encode<const P: u64>(m: &Rq<P, 1>) -> T64 { let delta = u64::MAX / P; // floored let coeffs = m.coeffs(); - Tn(array::from_fn(|i| T64(coeffs[i].0 * delta))) + // Tn(array::from_fn(|i| T64(coeffs[i].0 * delta))) + T64(coeffs[0].0 * delta) } - pub fn decode<const P: u64>(p: &Tn<1>) -> Rq<P, 1> { + pub fn decode<const P: u64>(p: &T64) -> 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(rng: impl Rng, sk: &SecretKey<K>, p: &Tn<1>) -> Result<Self> { - let glwe = GLWE::encrypt_s(rng, &sk.0, p)?; + pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<K>, p: &T64) -> Result<Self> { + let glwe = GLWE::encrypt_s(rng, &sk, p)?; Ok(Self(glwe)) } - pub fn encrypt(rng: impl Rng, pk: &PublicKey<K>, p: &Tn<1>) -> Result<Self> { - let glwe = GLWE::encrypt(rng, &pk.0, p)?; + pub fn encrypt(rng: impl Rng, pk: &PublicKey<K>, p: &T64) -> Result<Self> { + let glwe = GLWE::encrypt(rng, &pk, p)?; Ok(Self(glwe)) } - pub fn decrypt(&self, sk: &SecretKey<K>) -> Tn<1> { - self.0.decrypt(&sk.0) + pub fn decrypt(&self, sk: &SecretKey<K>) -> T64 { + self.0.decrypt(&sk) } }

@@ -86,29 +89,29 @@ }

} // plaintext addition -impl<const K: usize> Add<Tn<1>> for TLWE<K> { +impl<const K: usize> Add<T64> for TLWE<K> { type Output = Self; - fn add(self, plaintext: Tn<1>) -> Self { - let a: TR<Tn<1>, K> = self.0 .0; - let b: Tn<1> = self.0 .1 + plaintext; + fn add(self, plaintext: T64) -> Self { + let a: TR<T64, K> = self.0 .0; + let b: T64 = self.0 .1 + plaintext; Self(GLWE(a, b)) } } // plaintext substraction -impl<const K: usize> Sub<Tn<1>> for TLWE<K> { +impl<const K: usize> Sub<T64> for TLWE<K> { type Output = Self; - fn sub(self, plaintext: Tn<1>) -> Self { - let a: TR<Tn<1>, K> = self.0 .0; - let b: Tn<1> = self.0 .1 - plaintext; + fn sub(self, plaintext: T64) -> Self { + let a: TR<T64, K> = self.0 .0; + let b: T64 = self.0 .1 - plaintext; Self(GLWE(a, b)) } } // plaintext multiplication -impl<const K: usize> Mul<Tn<1>> for TLWE<K> { +impl<const K: usize> Mul<T64> for TLWE<K> { type Output = Self; - fn mul(self, plaintext: Tn<1>) -> Self { - let a: TR<Tn<1>, K> = TR(self.0 .0 .0.iter().map(|r_i| *r_i * plaintext).collect()); - let b: Tn<1> = self.0 .1 * plaintext; + fn mul(self, plaintext: T64) -> Self { + let a: TR<T64, K> = TR(self.0 .0 .0.iter().map(|r_i| *r_i * plaintext).collect()); + let b: T64 = self.0 .1 * plaintext; Self(GLWE(a, b)) } }

@@ -134,7 +137,7 @@ let (sk, pk) = S::new_key(&mut rng)?;

let m = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?; dbg!(&m); - let p: Tn<1> = S::encode::<T>(&m); + let p: T64 = S::encode::<T>(&m); dbg!(&p); let c = S::encrypt(&mut rng, &pk, &p)?;

@@ -168,8 +171,8 @@ let (sk, pk) = S::new_key(&mut rng)?;

let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?; let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?; - let p1: Tn<1> = S::encode::<T>(&m1); // plaintext - let p2: Tn<1> = S::encode::<T>(&m2); // plaintext + let p1: T64 = S::encode::<T>(&m1); // plaintext + let p2: T64 = S::encode::<T>(&m2); // plaintext let c1 = S::encrypt(&mut rng, &pk, &p1)?; let c2 = S::encrypt(&mut rng, &pk, &p2)?;

@@ -199,8 +202,8 @@ let (sk, pk) = S::new_key(&mut rng)?;

let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?; let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?; - let p1: Tn<1> = S::encode::<T>(&m1); // plaintext - let p2: Tn<1> = S::encode::<T>(&m2); // plaintext + let p1: T64 = S::encode::<T>(&m1); // plaintext + let p2: T64 = S::encode::<T>(&m2); // plaintext let c1 = S::encrypt(&mut rng, &pk, &p1)?;

@@ -209,7 +212,7 @@

let p3_recovered = c3.decrypt(&sk); let m3_recovered = S::decode::<T>(&p3_recovered); - assert_eq!((m1 + m2).remodule::<T>(), m3_recovered); + assert_eq!(m1 + m2, m3_recovered); } Ok(())

@@ -229,15 +232,16 @@ let (sk, pk) = S::new_key(&mut rng)?;

let m1 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?; let m2 = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?; - let p1: Tn<1> = S::encode::<T>(&m1); + let p1: T64 = S::encode::<T>(&m1); // don't scale up p2, set it directly from m2 - let p2: Tn<1> = Tn(array::from_fn(|i| T64(m2.coeffs()[i].0))); + // let p2: T64 = Tn(array::from_fn(|i| T64(m2.coeffs()[i].0))); + let p2: T64 = T64(m2.coeffs()[0].0); let c1 = S::encrypt(&mut rng, &pk, &p1)?; let c3 = c1 * p2; - let p3_recovered: Tn<1> = c3.decrypt(&sk); + let p3_recovered: T64 = c3.decrypt(&sk); let m3_recovered = S::decode::<T>(&p3_recovered); assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered); }