arnaucube/fhe-study
1
1
Fork 0
mirror of https://github.com/arnaucube/fhe-study.git synced 2026-01-23 20:23:54 +01:00
Code Issues Packages Projects Releases Wiki Activity

adapt gfhe to work with Ring trait, so that it can work with Rq & Tn (for TFHE)

Browse Source
This commit is contained in:
arnaucube
2025-07-24 14:26:48 +00:00
committed by arnaucube
parent 4790fdbb3b
commit 87da85a035
10 changed files with 428 additions and 185 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
arith/src/ring.rs
View File

@@ -3,7 +3,10 @@ use std::fmt::Debug;
use std::iter::Sum;
use std::ops::{Add, AddAssign, Mul, Sub, SubAssign};
/// Represents a ring element. Currently implemented by ring_n.rs#R and ring_nq.rs#Rq.
/// Represents a ring element. Currently implemented by ring_n.rs#R and
/// ring_nq.rs#Rq. Is not a 'pure algebraic ring', but more a custom trait
/// definition which includes methods like `mod_switch`.
// assumed to be mod (X^N +1)
pub trait Ring:
Sized
+ Add<Output = Self>
@@ -11,17 +14,22 @@ pub trait Ring:
+ Sum
+ Sub<Output = Self>
+ SubAssign
+ Mul<Output = Self>
+ Mul<u64, Output = Self> // scalar mul
+ Mul<Output = Self> // internal product
+ Mul<u64, Output = Self> // scalar mul, external product
+ Mul<Self::C, Output = Self>
+ PartialEq
+ Debug
+ Clone
+ Copy
+ Sum<<Self as Add>::Output>
+ Sum<<Self as Mul>::Output>
{
/// C defines the coefficient type
type C: Debug + Clone;
const Q: u64;
const N: usize;
fn coeffs(&self) -> Vec<Self::C>;
fn zero() -> Self;
// note/wip/warning: dist (0,q) with f64, will output more '0=q' elements than other values
@@ -31,6 +39,9 @@ pub trait Ring:
fn decompose(&self, beta: u32, l: u32) -> Vec<Self>;
fn remodule<const P: u64>(&self) -> impl Ring;
fn mod_switch<const P: u64>(&self) -> impl Ring;
/// 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

50
arith/src/ring_n.rs
View File

@@ -15,9 +15,13 @@ use crate::Ring;
#[derive(Clone, Copy)]
pub struct R<const N: usize>(pub [i64; N]);
impl<const N: usize> Ring for R<N> {
type C = i64;
fn coeffs(&self) -> Vec<Self::C> {
// impl<const N: usize> Ring for R<N> {
impl<const N: usize> R<N> {
// type C = i64;
// const Q: u64 = i64::MAX as u64; // WIP
// const N: usize = N;
pub fn coeffs(&self) -> Vec<i64> {
self.0.to_vec()
}
fn zero() -> Self {
@@ -32,30 +36,39 @@ impl<const N: usize> Ring for R<N> {
// Self(coeffs)
}
fn from_vec(coeffs: Vec<Self::C>) -> Self {
pub fn from_vec(coeffs: Vec<i64>) -> Self {
let mut p = coeffs;
modulus::<N>(&mut p);
Self(array::from_fn(|i| p[i]))
}
/*
// returns the decomposition of each polynomial coefficient
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 remodule<const P: u64>(&self) -> impl Ring {
unimplemented!()
}
fn mod_switch<const P: u64, const M: usize>(&self) -> R<N> {
unimplemented!()
}
/// 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)
// 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> {
@@ -65,9 +78,9 @@ impl<const Q: u64, const N: usize> From<crate::ring_nq::Rq<Q, N>> for R<N> {
}
impl<const N: usize> R<N> {
pub fn coeffs(&self) -> [i64; N] {
self.0
}
// pub fn coeffs(&self) -> [i64; N] {
// self.0
// }
pub fn to_rq<const Q: u64>(self) -> crate::Rq<Q, N> {
crate::Rq::<Q, N>::from(self)
}
@@ -318,6 +331,13 @@ pub fn mod_centered_q<const Q: u64, const N: usize>(p: Vec<i128>) -> R<N> {
R::<N>::from_vec(r.iter().map(|v| *v as i64).collect::<Vec<i64>>())
}
impl<const N: usize> Mul<i64> for R<N> {
type Output = Self;
fn mul(self, s: i64) -> Self {
self.mul_by_i64(s)
}
}
// mul by u64
impl<const N: usize> Mul<u64> for R<N> {
type Output = Self;

40
arith/src/ring_nq.rs
View File

@@ -30,6 +30,9 @@ pub struct Rq<const Q: u64, const N: usize> {
impl<const Q: u64, const N: usize> Ring for Rq<Q, N> {
type C = Zq<Q>;
const Q: u64 = Q;
const N: usize = N;
fn coeffs(&self) -> Vec<Self::C> {
self.coeffs.to_vec()
}
@@ -71,9 +74,26 @@ impl<const Q: u64, const N: usize> Ring for 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
// Warning: this method will behave differently depending on the values P and Q:
// if Q<P, it just 'renames' the modulus parameter to P
// if Q>=P, it crops to mod P
fn remodule<const P: u64>(&self) -> Rq<P, N> {
Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
}
/// perform the mod switch operation from Q to Q', where Q2=Q'
// fn mod_switch<const P: u64, const M: usize>(&self) -> impl Ring {
fn mod_switch<const P: u64>(&self) -> Rq<P, N> {
// assert_eq!(N, M); // sanity check
Rq::<P, N> {
coeffs: array::from_fn(|i| self.coeffs[i].mod_switch::<P>()),
evals: None,
}
}
/// 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()
@@ -183,17 +203,9 @@ impl<const Q: u64, const N: usize> Rq<Q, N> {
// Warning: this method will behave differently depending on the values P and Q:
// if Q<P, it just 'renames' the modulus parameter to P
// if Q>=P, it crops to mod P
pub fn remodule<const P: u64>(&self) -> Rq<P, N> {
Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
}
/// perform the mod switch operation from Q to Q', where Q2=Q'
pub fn mod_switch<const Q2: u64>(&self) -> Rq<Q2, N> {
Rq::<Q2, N> {
coeffs: array::from_fn(|i| self.coeffs[i].mod_switch::<Q2>()),
evals: None,
}
}
// pub fn remodule<const P: u64>(&self) -> Rq<P, N> {
// Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
// }
// applies mod(T) to all coefficients of self
pub fn coeffs_mod<const T: u64>(&self) -> Self {

30
arith/src/ring_torus.rs
View File

@@ -12,7 +12,7 @@ use std::array;
use std::iter::Sum;
use std::ops::{Add, AddAssign, Mul, Sub, SubAssign};
use crate::{ring::Ring, torus::T64};
use crate::{ring::Ring, torus::T64, Rq, Zq};
/// 𝕋_<N,Q>[X] = 𝕋<Q>[X]/(X^N +1), polynomials modulo X^N+1 with coefficients in
/// 𝕋, where Q=2^64.
@@ -22,6 +22,9 @@ pub struct Tn<const N: usize>(pub [T64; N]);
impl<const N: usize> Ring for Tn<N> {
type C = T64;
const Q: u64 = u64::MAX; // WIP
const N: usize = N;
fn coeffs(&self) -> Vec<T64> {
self.0.to_vec()
}
@@ -50,6 +53,22 @@ impl<const N: usize> Ring for Tn<N> {
r.iter().map(|a_i| Self::from_vec(a_i.clone())).collect()
}
fn remodule<const P: u64>(&self) -> Tn<N> {
todo!()
// Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
}
// fn mod_switch<const P: u64>(&self) -> impl Ring {
fn mod_switch<const P: u64>(&self) -> Rq<P, N> {
// unimplemented!()
// TODO WIP
let coeffs = array::from_fn(|i| Zq::<P>::from_u64(self.0[i].mod_switch::<P>().0));
Rq::<P, N> {
coeffs,
evals: None,
}
}
/// 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
@@ -174,12 +193,19 @@ fn modulus_u128<const N: usize>(p: &mut Vec<u128>) {
return;
}
for i in N..p.len() {
p[i - N] = p[i - N].clone() - p[i].clone();
// p[i - N] = p[i - N].clone() - p[i].clone();
p[i - N] = p[i - N].wrapping_sub(p[i]);
p[i] = 0;
}
p.truncate(N);
}
impl<const N: usize> Mul<T64> for Tn<N> {
type Output = Self;
fn mul(self, s: T64) -> Self {
Self(array::from_fn(|i| self.0[i] * s))
}
}
// mul by u64
impl<const N: usize> Mul<u64> for Tn<N> {
type Output = Self;

3
arith/src/torus.rs
View File

@@ -33,6 +33,9 @@ impl T64 {
.map(|i| T64(((self.0 >> i) & 1) as u64))
.collect()
}
pub fn mod_switch<const Q2: u64>(&self) -> T64 {
todo!()
}
}
impl Add<T64> for T64 {

50
gfhe/src/glev.rs
View File

@@ -3,7 +3,7 @@ use rand::Rng;
use rand_distr::{Normal, Uniform};
use std::ops::{Add, Mul};
use arith::{Ring, Rq, TR};
use arith::{Ring, TR};
use crate::glwe::{PublicKey, SecretKey, GLWE};
@@ -11,25 +11,19 @@ const ERR_SIGMA: f64 = 3.2;
// l GLWEs
#[derive(Clone, Debug)]
pub struct GLev<const Q: u64, const N: usize, const K: usize>(pub(crate) Vec<GLWE<Q, N, K>>);
pub struct GLev<R: Ring, const K: usize>(pub(crate) Vec<GLWE<R, 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>()
}
impl<R: Ring, const K: usize> GLev<R, K> {
pub fn encrypt(
mut rng: impl Rng,
beta: u32,
l: u32,
pk: &PublicKey<Q, N, K>,
m: &Rq<Q, N>,
pk: &PublicKey<R, K>,
m: &R,
) -> Result<Self> {
let glev: Vec<GLWE<Q, N, K>> = (1..l + 1)
let glev: Vec<GLWE<R, K>> = (0..l)
.map(|i| {
GLWE::<Q, N, K>::encrypt(&mut rng, pk, &(*m * (Q / beta.pow(i as u32) as u64)))
GLWE::<R, K>::encrypt(&mut rng, pk, &(*m * (R::Q / beta.pow(i as u32) as u64)))
})
.collect::<Result<Vec<_>>>()?;
@@ -39,21 +33,22 @@ impl<const Q: u64, const N: usize, const K: usize> GLev<Q, N, K> {
mut rng: impl Rng,
beta: u32,
l: u32,
sk: &SecretKey<Q, N, K>,
m: &Rq<Q, N>,
sk: &SecretKey<R, K>,
m: &R,
// delta: u64,
) -> Result<Self> {
let glev: Vec<GLWE<Q, N, K>> = (1..l + 1)
let glev: Vec<GLWE<R, K>> = (1..l + 1)
.map(|i| {
GLWE::<Q, N, K>::encrypt_s(&mut rng, sk, &(*m * (Q / beta.pow(i as u32) as u64)))
GLWE::<R, K>::encrypt_s(&mut rng, sk, &(*m * (R::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>, beta: u32) -> Rq<Q, N> {
let pt = self.0[0].decrypt(sk);
pt.mul_div_round(beta as u64, Q)
pub fn decrypt<const T: u64>(&self, sk: &SecretKey<R, K>, beta: u32) -> R {
let pt = self.0[1].decrypt(sk);
pt.mul_div_round(beta as u64, R::Q)
}
}
@@ -63,6 +58,7 @@ mod tests {
use rand::distributions::Uniform;
use super::*;
use arith::Rq;
#[test]
fn test_encrypt_decrypt() -> Result<()> {
@@ -70,25 +66,25 @@ mod tests {
const N: usize = 128;
const T: u64 = 2; // plaintext modulus
const K: usize = 16;
type S = GLev<Q, N, K>;
type S = GLev<Rq<Q, N>, K>;
let beta: u32 = 2;
let l: u32 = 16;
// let delta: u64 = Q / T; // floored
let mut rng = rand::thread_rng();
for _ in 0..200 {
let (sk, pk) = GLWE::<Q, N, K>::new_key(&mut rng)?;
let (sk, pk) = GLWE::<Rq<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 p: Rq<Q, N> = S::encode::<T>(&m); // plaintext
let m: Rq<Q, N> = m.remodule::<Q>();
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);
let c = S::encrypt(&mut rng, beta, l, &pk, &m)?;
let m_recovered = c.decrypt::<T>(&sk, beta);
assert_eq!(m, m_recovered);
assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
}
Ok(())

330
gfhe/src/glwe.rs
View File

@@ -1,3 +1,6 @@
//! Generalized LWE.
//!
use anyhow::Result;
use itertools::zip_eq;
use rand::Rng;
@@ -11,32 +14,34 @@ use crate::glev::GLev;
const ERR_SIGMA: f64 = 3.2;
/// GLWE implemented over the `Ring` trait, so that it can be also instantiated
/// over the Torus polynomials 𝕋_<N,q>[X] = 𝕋_q[X]/ (X^N+1).
#[derive(Clone, Debug)]
pub struct GLWE<const Q: u64, const N: usize, const K: usize>(TR<Rq<Q, N>, K>, Rq<Q, N>);
pub struct GLWE<R: Ring, const K: usize>(TR<R, K>, R);
#[derive(Clone, Debug)]
pub struct SecretKey<const Q: u64, const N: usize, const K: usize>(TR<Rq<Q, N>, K>);
pub struct SecretKey<R: Ring, const K: usize>(TR<R, K>);
#[derive(Clone, Debug)]
pub struct PublicKey<const Q: u64, const N: usize, const K: usize>(Rq<Q, N>, TR<Rq<Q, N>, K>);
pub struct PublicKey<R: Ring, const K: usize>(R, TR<R, K>);
// K GLevs, each KSK_i=l GLWEs
#[derive(Clone, Debug)]
pub struct KSK<const Q: u64, const N: usize, const K: usize>(Vec<GLev<Q, N, K>>);
pub struct KSK<R: Ring, const K: usize>(Vec<GLev<R, K>>);
impl<const Q: u64, const N: usize, const K: usize> GLWE<Q, N, K> {
impl<R: Ring, const K: usize> GLWE<R, K> {
pub fn zero() -> Self {
Self(TR::zero(), Rq::zero())
Self(TR::zero(), R::zero())
}
pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<Q, N, K>, PublicKey<Q, N, K>)> {
pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<R, K>, PublicKey<R, K>)> {
let Xi_key = Uniform::new(0_f64, 2_f64);
let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
let s: TR<Rq<Q, N>, K> = TR::rand(&mut rng, Xi_key);
let a: TR<Rq<Q, N>, K> = TR::rand(&mut rng, Uniform::new(0_f64, Q as f64));
let e = Rq::<Q, N>::rand(&mut rng, Xi_err);
let s: TR<R, K> = TR::rand(&mut rng, Xi_key);
let a: TR<R, K> = TR::rand(&mut rng, Uniform::new(0_f64, R::Q as f64));
let e = R::rand(&mut rng, Xi_err);
let pk: PublicKey<Q, N, K> = PublicKey((&a * &s) + e, a);
let pk: PublicKey<R, K> = PublicKey((&a * &s) + e, a);
Ok((SecretKey(s), pk))
}
@@ -44,42 +49,85 @@ impl<const Q: u64, const N: usize, const K: usize> GLWE<Q, N, K> {
mut rng: impl Rng,
beta: u32,
l: u32,
sk: &SecretKey<Q, N, K>,
new_sk: &SecretKey<Q, N, K>,
) -> Result<KSK<Q, N, K>> {
let r: Vec<GLev<Q, N, K>> = (0..K)
sk: &SecretKey<R, K>,
new_sk: &SecretKey<R, K>,
) -> Result<KSK<R, K>> {
let r: Vec<GLev<R, K>> = (0..K)
.into_iter()
.map(|i|
// treat sk_i as the msg being encrypted
GLev::<Q, N, K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0 .0[i]))
GLev::<R, K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0 .0[i]))
.collect::<Result<Vec<_>>>()?;
Ok(KSK(r))
}
pub fn key_switch(&self, beta: u32, l: u32, ksk: &KSK<Q, N, K>) -> Self {
let (a, b): (TR<Rq<Q, N>, K>, Rq<Q, N>) = (self.0.clone(), self.1);
pub fn key_switch(&self, beta: u32, l: u32, ksk: &KSK<R, K>) -> Self {
let (a, b): (TR<R, K>, R) = (self.0.clone(), self.1);
let lhs: GLWE<Q, N, K> = GLWE(TR::zero(), b);
let lhs: GLWE<R, K> = GLWE(TR::zero(), b);
// K iterations, ksk.0 contains K times GLev
let rhs: GLWE<Q, N, K> = zip_eq(a.0, ksk.0.clone())
let rhs: GLWE<R, K> = zip_eq(a.0, ksk.0.clone())
.map(|(a_i, ksk_i)| Self::dot_prod(a_i.decompose(beta, l), ksk_i))
.sum();
lhs - rhs
}
// note: a_decomp is of length N
fn dot_prod(a_decomp: Vec<Rq<Q, N>>, ksk_i: GLev<Q, N, K>) -> GLWE<Q, N, K> {
fn dot_prod(a_decomp: Vec<R>, ksk_i: GLev<R, K>) -> GLWE<R, K> {
// l times GLWES
let glwes: Vec<GLWE<Q, N, K>> = ksk_i.0;
let glwes: Vec<GLWE<R, K>> = ksk_i.0;
// l iterations
let r: GLWE<Q, N, K> = zip_eq(a_decomp, glwes)
let r: GLWE<R, K> = zip_eq(a_decomp, glwes)
.map(|(a_d_i, glwe_i)| glwe_i * a_d_i)
.sum();
r
}
// encrypts with the given SecretKey (instead of PublicKey)
pub fn encrypt_s(
mut rng: impl Rng,
sk: &SecretKey<R, K>,
m: &R, // already scaled
) -> Result<Self> {
let Xi_key = Uniform::new(0_f64, 2_f64);
let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
let a: TR<R, K> = TR::rand(&mut rng, Xi_key);
let e = R::rand(&mut rng, Xi_err);
let b: R = (&a * &sk.0) + *m + e;
Ok(Self(a, b))
}
pub fn encrypt(
mut rng: impl Rng,
pk: &PublicKey<R, K>,
m: &R, // already scaled
) -> Result<Self> {
let Xi_key = Uniform::new(0_f64, 2_f64);
let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
let u: R = R::rand(&mut rng, Xi_key);
let e0 = R::rand(&mut rng, Xi_err);
let e1 = TR::<R, K>::rand(&mut rng, Xi_err);
let b: R = pk.0.clone() * u.clone() + *m + e0;
let d: TR<R, K> = &pk.1 * &u + e1;
Ok(Self(d, b))
}
// returns m' not downscaled
pub fn decrypt(&self, sk: &SecretKey<R, K>) -> R {
let (d, b): (TR<R, K>, R) = (self.0.clone(), self.1);
let p: R = b - &d * &sk.0;
p
}
}
// Methods for when Ring=Rq<Q,N>
impl<const Q: u64, const N: usize, const K: usize> GLWE<Rq<Q, N>, K> {
// scale up
pub fn encode<const T: u64>(m: &Rq<T, N>) -> Rq<Q, N> {
let m = m.remodule::<Q>();
@@ -87,80 +135,54 @@ impl<const Q: u64, const N: usize, const K: usize> GLWE<Q, N, K> {
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>) -> 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 + e;
Ok(Self(a, b))
}
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)?;
let u: Rq<Q, N> = Rq::rand(&mut rng, Xi_key);
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 + e0;
let d: TR<Rq<Q, N>, K> = &pk.1 * &u + e1;
Ok(Self(d, b))
}
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;
pub fn decode<const T: u64>(m: &Rq<Q, N>) -> Rq<T, N> {
let r = m.mul_div_round(T, Q);
let r: Rq<T, N> = r.remodule::<T>();
r
}
pub fn mod_switch<const P: u64>(&self) -> GLWE<P, N, K> {
let a: TR<Rq<P, N>, K> = TR(self.0 .0.iter().map(|r| r.mod_switch::<P>()).collect());
pub fn mod_switch<const P: u64>(&self) -> GLWE<Rq<P, N>, K> {
let a: TR<Rq<P, N>, K> = TR(self
.0
.0
.iter()
.map(|r| r.mod_switch::<P>())
.collect::<Vec<_>>());
let b: Rq<P, N> = self.1.mod_switch::<P>();
GLWE(a, b)
}
}
impl<const Q: u64, const N: usize, const K: usize> Add<GLWE<Q, N, K>> for GLWE<Q, N, K> {
impl<R: Ring, const K: usize> Add<GLWE<R, K>> for GLWE<R, K> {
type Output = Self;
fn add(self, other: Self) -> Self {
let a: TR<Rq<Q, N>, K> = self.0 + other.0;
let b: Rq<Q, N> = self.1 + other.1;
let a: TR<R, K> = self.0 + other.0;
let b: R = self.1 + other.1;
Self(a, b)
}
}
impl<const Q: u64, const N: usize, const K: usize> Add<Rq<Q, N>> for GLWE<Q, N, K> {
impl<R: Ring, const K: usize> Add<R> for GLWE<R, K> {
type Output = Self;
fn add(self, plaintext: Rq<Q, N>) -> Self {
let a: TR<Rq<Q, N>, K> = self.0;
let b: Rq<Q, N> = self.1 + plaintext;
fn add(self, plaintext: R) -> Self {
let a: TR<R, K> = self.0;
let b: R = self.1 + plaintext;
Self(a, b)
}
}
impl<const Q: u64, const N: usize, const K: usize> AddAssign for GLWE<Q, N, K> {
impl<R: Ring, const K: usize> AddAssign for GLWE<R, K> {
fn add_assign(&mut self, rhs: Self) {
for i in 0..K {
self.0 .0[i] = self.0 .0[i] + rhs.0 .0[i];
self.0 .0[i] = self.0 .0[i].clone() + rhs.0 .0[i].clone();
}
self.1 = self.1 + rhs.1;
self.1 = self.1.clone() + rhs.1.clone();
}
}
impl<const Q: u64, const N: usize, const K: usize> Sum<GLWE<Q, N, K>> for GLWE<Q, N, K> {
impl<R: Ring, const K: usize> Sum<GLWE<R, K>> for GLWE<R, K> {
fn sum<I>(iter: I) -> Self
where
I: Iterator<Item = Self>,
{
let mut acc = GLWE::<Q, N, K>::zero();
let mut acc = GLWE::<R, K>::zero();
for e in iter {
acc += e;
}
@@ -168,37 +190,60 @@ impl<const Q: u64, const N: usize, const K: usize> Sum<GLWE<Q, N, K>> for GLWE<Q
}
}
impl<const Q: u64, const N: usize, const K: usize> Sub<GLWE<Q, N, K>> for GLWE<Q, N, K> {
impl<R: Ring, const K: usize> Sub<GLWE<R, K>> for GLWE<R, K> {
type Output = Self;
fn sub(self, other: Self) -> Self {
let a: TR<Rq<Q, N>, K> = self.0 - other.0;
let b: Rq<Q, N> = self.1 - other.1;
let a: TR<R, K> = self.0 - other.0;
let b: R = self.1 - other.1;
Self(a, b)
}
}
impl<const Q: u64, const N: usize, const K: usize> Mul<Rq<Q, N>> for GLWE<Q, N, K> {
impl<R: Ring, const K: usize> Mul<R> for GLWE<R, K> {
type Output = Self;
fn mul(self, plaintext: Rq<Q, N>) -> Self {
// first compute the NTT for plaintext, to avoid computing it at each
// iteration, speeding up the multiplications
let mut plaintext = plaintext.clone();
plaintext.compute_evals();
let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
let b: Rq<Q, N> = self.1 * plaintext;
fn mul(self, plaintext: R) -> Self {
let a: TR<R, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
let b: R = self.1 * plaintext;
Self(a, b)
}
}
// for when R = Rq<Q,N>
// impl<const Q: u64, const N: usize, const K: usize> Mul<Rq<Q, N>> for GLWE<Rq<Q, N>, K> {
// type Output = Self;
// fn mul(self, plaintext: Rq<Q, N>) -> Self {
// // first compute the NTT for plaintext, to avoid computing it at each
// // iteration, speeding up the multiplications
// let mut plaintext = plaintext.clone();
// plaintext.compute_evals();
//
// let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect());
// let b: Rq<Q, N> = self.1 * plaintext;
// Self(a, b)
// }
// }
impl<const Q: u64, const N: usize, const K: usize> Mul<Zq<Q>> for GLWE<Q, N, K> {
type Output = Self;
fn mul(self, e: Zq<Q>) -> Self {
let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e).collect());
let b: Rq<Q, N> = self.1 * e;
Self(a, b)
}
}
// impl<R: Ring, const K: usize> Mul<R::C> for GLWE<R, K>
// // where
// // // R: std::ops::Mul<<R as arith::Ring>::C>,
// // // Vec<R>: FromIterator<<R as Mul<<R as arith::Ring>::C>>::Output>,
// // Vec<R>: FromIterator<<R as Mul<<R as arith::Ring>::C>>::Output>,
// {
// type Output = Self;
// fn mul(self, e: R::C) -> Self {
// let a: TR<R, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e.clone()).collect());
// let b: R = self.1 * e.clone();
// Self(a, b)
// }
// }
// impl<const Q: u64, const N: usize, const K: usize> Mul<Zq<Q>> for GLWE<Q, N, K> {
// type Output = Self;
// fn mul(self, e: Zq<Q>) -> Self {
// let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e).collect());
// let b: Rq<Q, N> = self.1 * e;
// Self(a, b)
// }
// }
#[cfg(test)]
mod tests {
@@ -213,7 +258,7 @@ mod tests {
const N: usize = 128;
const T: u64 = 32; // plaintext modulus
const K: usize = 16;
type S = GLWE<Q, N, K>;
type S = GLWE<Rq<Q, N>, K>;
let mut rng = rand::thread_rng();
@@ -221,10 +266,11 @@ mod tests {
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 p = S::encode::<T>(&m); // plaintext
let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; // msg
// let m: Rq<Q, N> = m.remodule::<Q>();
let c = S::encrypt(&mut rng, &pk, &p)?;
let p = S::encode::<T>(&m); // plaintext
let c = S::encrypt(&mut rng, &pk, &p)?; // ciphertext
let p_recovered = c.decrypt(&sk);
let m_recovered = S::decode::<T>(&p_recovered);
@@ -241,13 +287,57 @@ mod tests {
Ok(())
}
use arith::{Tn, T64};
use std::array;
pub fn t_encode<const P: u64>(m: &Rq<P, 4>) -> Tn<4> {
let delta = u64::MAX / P; // floored
let coeffs = m.coeffs();
Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
}
pub fn t_decode<const P: u64>(p: &Tn<4>) -> Rq<P, 4> {
let p = p.mul_div_round(P, u64::MAX);
Rq::<P, 4>::from_vec_u64(p.coeffs().iter().map(|c| c.0).collect())
}
#[test]
fn test_encrypt_decrypt_torus() -> Result<()> {
const N: usize = 128;
const T: u64 = 32; // plaintext modulus
const K: usize = 16;
type S = GLWE<Tn<4>, 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_f64, T as f64);
let m = Rq::<T, 4>::rand(&mut rng, msg_dist); // msg
let p = t_encode::<T>(&m); // plaintext
let c = S::encrypt(&mut rng, &pk, &p)?; // ciphertext
let p_recovered = c.decrypt(&sk);
let m_recovered = t_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 = t_decode::<T>(&p_recovered);
assert_eq!(m, m_recovered);
}
Ok(())
}
#[test]
fn test_addition() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 128;
const T: u64 = 20;
const K: usize = 16;
type S = GLWE<Q, N, K>;
type S = GLWE<Rq<Q, N>, K>;
let mut rng = rand::thread_rng();
@@ -280,7 +370,7 @@ mod tests {
const N: usize = 128;
const T: u64 = 32;
const K: usize = 16;
type S = GLWE<Q, N, K>;
type S = GLWE<Rq<Q, N>, K>;
let mut rng = rand::thread_rng();
@@ -300,7 +390,7 @@ mod tests {
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).remodule::<T>(), m3_recovered.remodule::<T>());
}
Ok(())
@@ -312,7 +402,7 @@ mod tests {
const N: usize = 16;
const T: u64 = 4;
const K: usize = 16;
type S = GLWE<Q, N, K>;
type S = GLWE<Rq<Q, N>, K>;
let mut rng = rand::thread_rng();
@@ -323,14 +413,14 @@ mod tests {
let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let p1: Rq<Q, N> = S::encode::<T>(&m1); // plaintext
let p2: Rq<Q, N> = m2.remodule::<Q>();
let p2 = m2.remodule::<Q>(); // notice we don't encode (scale by delta)
let c1 = S::encrypt(&mut rng, &pk, &p1)?;
let c3 = c1 * p2;
let p3_recovered: Rq<Q, N> = c3.decrypt(&sk);
let m3_recovered = S::decode::<T>(&p3_recovered);
let m3_recovered: Rq<T, N> = S::decode::<T>(&p3_recovered);
assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered);
}
@@ -343,35 +433,27 @@ mod tests {
const P: u64 = 2u64.pow(8) + 1;
// note: wip, Q and P chosen so that P/Q is an integer
const N: usize = 8;
const T: u64 = 8; // plaintext modulus, must be a prime or power of a prime
const T: u64 = 4; // plaintext modulus, must be a prime or power of a prime
const K: usize = 16;
type S = GLWE<Q, N, K>;
type S = GLWE<Rq<Q, N>, K>;
let delta: u64 = Q / T; // floored
let mut rng = rand::thread_rng();
dbg!(P as f64 / Q as f64);
dbg!(delta);
dbg!(delta as f64 * P as f64 / Q as f64);
dbg!(delta as f64 * (P as f64 / Q as f64));
for _ in 0..200 {
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 p = S::encode::<T>(&m); // plaintext
let p = S::encode::<T>(&m);
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> =
let c2: GLWE<Rq<P, N>, K> = c.mod_switch::<P>();
let sk2: SecretKey<Rq<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 p_recovered = c2.decrypt(&sk2);
let m_recovered = GLWE::<P, N, K>::decode::<T>(&p_recovered);
let m_recovered = GLWE::<Rq<P, N>, K>::decode::<T>(&p_recovered);
assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
}
@@ -385,7 +467,7 @@ mod tests {
const N: usize = 128;
const T: u64 = 2; // plaintext modulus
const K: usize = 16;
type S = GLWE<Q, N, K>;
type S = GLWE<Rq<Q, N>, K>;
let beta: u32 = 2;
let l: u32 = 16;
@@ -399,8 +481,8 @@ mod tests {
let msg_dist = Uniform::new(0_u64, T);
let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let p: Rq<Q, N> = S::encode::<T>(&m); // plaintext
let p = S::encode::<T>(&m); // plaintext
//
let c = S::encrypt_s(&mut rng, &sk, &p)?;
let c2 = c.key_switch(beta, l, &ksk);
@@ -408,14 +490,14 @@ mod tests {
// decrypt with the 2nd secret key
let p_recovered = c2.decrypt(&sk2);
let m_recovered = S::decode::<T>(&p_recovered);
assert_eq!(m, m_recovered);
assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>());
// do the same but now encrypting with pk
// let c = S::encrypt(&mut rng, &pk, &p)?;
// let c2 = c.key_switch(beta, l, &ksk);
// let p_recovered = c2.decrypt(&sk2);
// let m_recovered = S::decode::<T>(&p_recovered);
// assert_eq!(m, m_recovered);
let c = S::encrypt(&mut rng, &pk, &p)?;
let c2 = c.key_switch(beta, l, &ksk);
let p_recovered = c2.decrypt(&sk2);
let m_recovered = S::decode::<T>(&p_recovered);
assert_eq!(m, m_recovered);
Ok(())
}

3
gfhe/src/lib.rs
View File

@@ -7,3 +7,6 @@
pub mod glev;
pub mod glwe;
pub use glev::GLev;
pub use glwe::GLWE;

1
tfhe/Cargo.toml
View File

@@ -10,3 +10,4 @@ rand_distr = { workspace = true }
itertools = { workspace = true }
arith = { path="../arith" }
gfhe = { path="../gfhe" }

89
tfhe/src/lib.rs
View File

@@ -5,5 +5,94 @@
#![allow(clippy::upper_case_acronyms)]
#![allow(dead_code)] // TMP
use anyhow::Result;
use rand::Rng;
use rand_distr::{Normal, Uniform};
use std::array;
use arith::{Ring, Rq, Tn, T64};
use gfhe::{glwe, GLWE};
pub mod tlev;
pub mod tlwe;
#[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 TLWE<const K: usize>(pub GLWE<Tn<1>, K>);
impl<const K: usize> TLWE<K> {
pub fn new_key(rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> {
let (sk, pk) = GLWE::new_key(rng)?;
Ok((SecretKey(sk), PublicKey(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())
}
pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<K>, p: &Tn<1>) -> Result<Self> {
let glwe = GLWE::encrypt_s(rng, &sk.0, 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)?;
Ok(Self(glwe))
}
pub fn decrypt(&self, sk: &SecretKey<K>) -> Tn<1> {
self.0.decrypt(&sk.0)
}
}
#[cfg(test)]
mod tests {
use anyhow::Result;
use rand::distributions::Uniform;
use super::*;
#[test]
fn test_encrypt_decrypt() -> Result<()> {
const T: u64 = 128; // plaintext modulus
const K: usize = 16;
type S = TLWE<K>;
// let delta: u64 = Q / T; // floored
let mut rng = rand::thread_rng();
for _ in 0..200 {
let (sk, pk) = S::new_key(&mut rng)?;
let msg_dist = Uniform::new(0_f64, T as f64);
let m = Rq::<T, 1>::rand(&mut rng, msg_dist); // msg
// let m: Rq<Q, N> = m.remodule::<Q>();
let p = S::encode::<T>(&m); // plaintext
let c = S::encrypt(&mut rng, &pk, &p)?; // ciphertext
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.remodule::<T>(), m_recovered.remodule::<T>());
}
Ok(())
}
}