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

tfhe: add TGLWE.sample_extraction

Browse Source
This commit is contained in:
arnaucube
2025-08-02 20:27:36 +00:00
parent 752525a1c6
commit 0617b8b9ff
4 changed files with 113 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

90
tfhe/src/tglwe.rs
View File

@@ -11,8 +11,23 @@ use arith::{Ring, Rq, Tn, T64, TR};
use gfhe::{glwe, GLWE};
use crate::tlev::TLev;
use crate::{tlwe, tlwe::TLWE};
// pub type SecretKey<const N: usize, const K: usize> = glwe::SecretKey<Tn<N>, K>;
#[derive(Clone, Debug)]
pub struct SecretKey<const N: usize, const K: usize>(pub glwe::SecretKey<Tn<N>, K>);
// pub struct SecretKey<const K: usize>(pub tlwe::SecretKey<K>);
impl<const N: usize, const K: usize> SecretKey<N, K> {
pub fn to_tlwe<const KN: usize>(self) -> tlwe::SecretKey<K> {
let s: TR<Tn<N>, K> = self.0 .0;
let r: Vec<Vec<T64>> = s.0.iter().map(|s_i| s_i.coeffs()).collect();
let r: Vec<T64> = r.into_iter().flatten().collect();
tlwe::SecretKey(glwe::SecretKey(TR(r)))
}
}
pub type SecretKey<const N: usize, const K: usize> = glwe::SecretKey<Tn<N>, K>;
pub type PublicKey<const N: usize, const K: usize> = glwe::PublicKey<Tn<N>, K>;
#[derive(Clone, Debug)]
@@ -23,9 +38,14 @@ impl<const N: usize, const K: usize> TGLWE<N, K> {
Self(GLWE::<Tn<N>, K>::zero())
}
pub fn new_key(rng: impl Rng) -> Result<(SecretKey<N, K>, PublicKey<N, K>)> {
let (sk, pk) = GLWE::new_key(rng)?;
// Ok((SecretKey(sk), PublicKey(pk)))
pub fn new_key<const KN: usize>(
mut rng: impl Rng,
) -> Result<(SecretKey<N, K>, PublicKey<N, K>)> {
assert_eq!(KN, K * N); // this is wip, while not being able to compute K*N
let (sk_tlwe, _) = TLWE::<KN>::new_key(&mut rng)?;
// let sk = crate::tlwe::sk_to_tglwe::<N, K, KN>(sk_tlwe);
let sk = sk_tlwe.to_tglwe::<N, K>();
let pk = GLWE::pk_from_sk(rng, sk.0.clone())?;
Ok((sk, pk))
}
@@ -41,7 +61,7 @@ impl<const N: usize, const K: usize> TGLWE<N, K> {
// encrypts with the given SecretKey (instead of PublicKey)
pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<N, K>, p: &Tn<N>) -> Result<Self> {
let glwe = GLWE::encrypt_s(rng, &sk, p)?;
let glwe = GLWE::encrypt_s(rng, &sk.0, p)?;
Ok(Self(glwe))
}
pub fn encrypt(rng: impl Rng, pk: &PublicKey<N, K>, p: &Tn<N>) -> Result<Self> {
@@ -49,7 +69,27 @@ impl<const N: usize, const K: usize> TGLWE<N, K> {
Ok(Self(glwe))
}
pub fn decrypt(&self, sk: &SecretKey<N, K>) -> Tn<N> {
self.0.decrypt(&sk)
self.0.decrypt(&sk.0)
}
/// Sample extraction / Coefficient extraction
pub fn sample_extraction(&self, h: usize) -> TLWE<K> {
assert!(h < N);
let a: TR<Tn<N>, K> = self.0 .0.clone();
// set a_{n*i+j} = a_{i, h-j} if j \in {0, h}
// -a_{i, n+h-j} if j \in {h+1, n-1}
let new_a: Vec<T64> = a
.iter()
.flat_map(|a_i| {
let a_i = a_i.coeffs();
(0..N)
.map(|j| if j <= h { a_i[h - j] } else { -a_i[N + h - j] })
.collect::<Vec<T64>>()
})
.collect::<Vec<T64>>();
TLWE(GLWE(TR(new_a), self.0 .1.coeffs()[h]))
}
}
@@ -130,7 +170,7 @@ mod tests {
let msg_dist = Uniform::new(0_u64, T);
for _ in 0..200 {
let (sk, pk) = S::new_key(&mut rng)?;
let (sk, pk) = TGLWE::<N, K>::new_key::<{ K * N }>(&mut rng)?;
let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let p: Tn<N> = S::encode::<T>(&m);
@@ -163,7 +203,7 @@ mod tests {
let msg_dist = Uniform::new(0_u64, T);
for _ in 0..200 {
let (sk, pk) = S::new_key(&mut rng)?;
let (sk, pk) = S::new_key::<{ K * N }>(&mut rng)?;
let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
@@ -195,7 +235,7 @@ mod tests {
let msg_dist = Uniform::new(0_u64, T);
for _ in 0..200 {
let (sk, pk) = S::new_key(&mut rng)?;
let (sk, pk) = S::new_key::<{ K * N }>(&mut rng)?;
let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
@@ -226,7 +266,7 @@ mod tests {
let msg_dist = Uniform::new(0_u64, T);
for _ in 0..200 {
let (sk, pk) = S::new_key(&mut rng)?;
let (sk, pk) = S::new_key::<{ K * N }>(&mut rng)?;
let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
@@ -245,4 +285,34 @@ mod tests {
Ok(())
}
#[test]
fn test_sample_extraction() -> Result<()> {
const T: u64 = 128; // msg space (msg modulus)
const N: usize = 64;
const K: usize = 16;
let mut rng = rand::thread_rng();
let msg_dist = Uniform::new(0_u64, T);
for _ in 0..20 {
let (sk, pk) = TGLWE::<N, K>::new_key::<{ K * N }>(&mut rng)?;
let sk_tlwe = sk.to_tlwe::<{ K * N }>();
let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let p: Tn<N> = TGLWE::<N, K>::encode::<T>(&m);
let c = TGLWE::<N, K>::encrypt(&mut rng, &pk, &p)?;
for h in 0..N {
let c_h: TLWE<K> = c.sample_extraction(h);
let p_recovered = c_h.decrypt(&sk_tlwe);
let m_recovered = TLWE::<K>::decode::<T>(&p_recovered);
assert_eq!(m.coeffs()[h], m_recovered.coeffs()[0]);
}
}
Ok(())
}
}

2
tfhe/src/tgsw.rs
View File

@@ -23,7 +23,7 @@ impl<const K: usize> TGSW<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)))
.map(|i| TLev::encrypt_s(&mut rng, beta, l, sk, &(-sk.0 .0 .0[i] * *m)))
.collect::<Result<Vec<_>>>()?;
let b: TLev<K> = TLev::encrypt_s(&mut rng, beta, l, sk, m)?;
Ok(Self(a, b))

30
tfhe/src/tlwe.rs
View File

@@ -13,8 +13,24 @@ use gfhe::{glwe, GLWE};
use crate::tlev::TLev;
// #[derive(Clone, Debug)]
// pub struct SecretKey<const K: usize>(glwe::SecretKey<T64, K>);
pub type SecretKey<const K: usize> = glwe::SecretKey<T64, K>;
pub struct SecretKey<const K: usize>(pub glwe::SecretKey<T64, K>);
// pub type SecretKey<const K: usize> = glwe::SecretKey<T64, K>;
impl<const KN: usize> SecretKey<KN> {
/// from TFHE [2018-421] paper: A TLWE key k \in B^n, can be interpreted as a
/// TRLWE key K \in B_N[X]^k having the same sequence of coefficients and
/// vice-versa.
pub fn to_tglwe<const N: usize, const K: usize>(self) -> crate::tglwe::SecretKey<N, K> {
let s: TR<T64, KN> = self.0 .0;
// split into K vectors, and interpret each of them as a T_N[X]/(X^N+1)
// polynomial
let r: Vec<Tn<N>> =
s.0.chunks(N)
.map(|v| Tn::<N>::from_vec(v.to_vec()))
.collect();
crate::tglwe::SecretKey(glwe::SecretKey::<Tn<N>, K>(TR(r)))
}
}
// #[derive(Clone, Debug)]
// pub struct PublicKey<const K: usize>(glwe::PublicKey<T64, K>);
@@ -32,9 +48,9 @@ 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)?;
let (sk, pk): (glwe::SecretKey<T64, K>, glwe::PublicKey<T64, K>) = GLWE::new_key(rng)?;
// Ok((SecretKey(sk), PublicKey(pk)))
Ok((sk, pk))
Ok((SecretKey(sk), pk))
}
pub fn encode<const P: u64>(m: &Rq<P, 1>) -> T64 {
@@ -50,7 +66,7 @@ impl<const K: usize> TLWE<K> {
// encrypts with the given SecretKey (instead of PublicKey)
pub fn encrypt_s(rng: impl Rng, sk: &SecretKey<K>, p: &T64) -> Result<Self> {
let glwe = GLWE::encrypt_s(rng, &sk, p)?;
let glwe = GLWE::encrypt_s(rng, &sk.0, p)?;
Ok(Self(glwe))
}
pub fn encrypt(rng: impl Rng, pk: &PublicKey<K>, p: &T64) -> Result<Self> {
@@ -58,7 +74,7 @@ impl<const K: usize> TLWE<K> {
Ok(Self(glwe))
}
pub fn decrypt(&self, sk: &SecretKey<K>) -> T64 {
self.0.decrypt(&sk)
self.0.decrypt(&sk.0)
}
pub fn new_ksk(
@@ -72,7 +88,7 @@ impl<const K: usize> TLWE<K> {
.into_iter()
.map(|i|
// treat sk_i as the msg being encrypted
TLev::<K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0 .0[i]))
TLev::<K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0.0 .0[i]))
.collect::<Result<Vec<_>>>()?;
Ok(KSK(r))