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 blind_rotation & lookup table computation

Browse Source
This commit is contained in:
arnaucube
2025-08-03 19:12:47 +00:00
parent 7bfcf6f7c1
commit 2c20a2ed0e
5 changed files with 101 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
tfhe/src/tglwe.rs
View File

@@ -19,12 +19,12 @@ pub struct SecretKey<const N: usize, const K: usize>(pub glwe::SecretKey<Tn<N>,
// 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;
pub fn to_tlwe<const KN: usize>(&self) -> tlwe::SecretKey<KN> {
let s: TR<Tn<N>, K> = self.0 .0.clone();
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)))
tlwe::SecretKey::<KN>(glwe::SecretKey::<T64, KN>(TR::<T64, KN>::new(r)))
}
}
@@ -37,15 +37,18 @@ impl<const N: usize, const K: usize> TGLWE<N, K> {
pub fn zero() -> Self {
Self(GLWE::<Tn<N>, K>::zero())
}
pub fn from_plaintext(p: Tn<N>) -> Self {
Self(GLWE::<Tn<N>, K>::from_plaintext(p))
}
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
// 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())?;
let pk: PublicKey<N, K> = GLWE::pk_from_sk(rng, sk.0.clone())?;
Ok((sk, pk))
}
@@ -73,7 +76,7 @@ impl<const N: usize, const K: usize> TGLWE<N, K> {
}
/// Sample extraction / Coefficient extraction
pub fn sample_extraction(&self, h: usize) -> TLWE<K> {
pub fn sample_extraction<const KN: usize>(&self, h: usize) -> TLWE<KN> {
assert!(h < N);
let a: TR<Tn<N>, K> = self.0 .0.clone();
@@ -91,6 +94,11 @@ impl<const N: usize, const K: usize> TGLWE<N, K> {
TLWE(GLWE(TR(new_a), self.0 .1.coeffs()[h]))
}
pub fn left_rotate(&self, h: usize) -> Self {
dbg!(&h);
let (a, b): (TR<Tn<N>, K>, Tn<N>) = (self.0 .0.clone(), self.0 .1);
Self(GLWE(a.left_rotate(h), b.left_rotate(h)))
}
}
impl<const N: usize, const K: usize> Add<TGLWE<N, K>> for TGLWE<N, K> {
@@ -291,13 +299,14 @@ mod tests {
const T: u64 = 128; // msg space (msg modulus)
const N: usize = 64;
const K: usize = 16;
const KN: usize = K * N;
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 (sk, pk) = TGLWE::<N, K>::new_key::<KN>(&mut rng)?;
let sk_tlwe = sk.to_tlwe::<KN>();
let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let p: Tn<N> = TGLWE::<N, K>::encode::<T>(&m);
@@ -305,10 +314,10 @@ mod tests {
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 c_h: TLWE<KN> = c.sample_extraction(h);
let p_recovered = c_h.decrypt(&sk_tlwe);
let m_recovered = TLWE::<K>::decode::<T>(&p_recovered);
let m_recovered = TLWE::<KN>::decode::<T>(&p_recovered);
assert_eq!(m.coeffs()[h], m_recovered.coeffs()[0]);
}
}

5
tfhe/src/tgsw.rs
View File

@@ -7,7 +7,10 @@ use std::ops::{Add, Mul};
use arith::{Ring, Rq, Tn, T64, TR};
use crate::tlev::TLev;
use crate::tlwe::{PublicKey, SecretKey, TLWE};
use crate::{
tglwe::TGLWE,
tlwe::{PublicKey, SecretKey, TLWE},
};
use gfhe::glwe::GLWE;
/// vector of length K+1 = [K], [1]

1
tfhe/src/tlev.rs
View File

@@ -6,6 +6,7 @@ use std::ops::{Add, Mul};
use arith::{Ring, Rq, Tn, T64, TR};
use crate::tglwe::TGLWE;
use crate::tlwe::{PublicKey, SecretKey, TLWE};
#[derive(Clone, Debug)]

86
tfhe/src/tlwe.rs
View File

@@ -1,20 +1,17 @@
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, T64, TR};
use arith::{Ring, Rq, Tn, Zq, T64, TR};
use gfhe::{glwe, GLWE};
use crate::tggsw::TGGSW;
use crate::tlev::TLev;
use crate::{tglwe, tglwe::TGLWE};
// #[derive(Clone, Debug)]
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
@@ -32,8 +29,6 @@ impl<const KN: usize> SecretKey<KN> {
}
}
// #[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)]
@@ -49,14 +44,12 @@ impl<const K: usize> TLWE<K> {
pub fn new_key(rng: impl Rng) -> Result<(SecretKey<K>, PublicKey<K>)> {
let (sk, pk): (glwe::SecretKey<T64, K>, glwe::PublicKey<T64, K>) = GLWE::new_key(rng)?;
// Ok((SecretKey(sk), PublicKey(pk)))
Ok((SecretKey(sk), pk))
}
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)))
T64(coeffs[0].0 * delta)
}
pub fn decode<const P: u64>(p: &T64) -> Rq<P, 1> {
@@ -105,6 +98,76 @@ impl<const K: usize> TLWE<K> {
lhs - rhs
}
// modulus switch from Q (2^64) to Q2 (in blind_rotation Q2=K*N)
pub fn mod_switch<const Q2: u64>(&self) -> Self {
let a: TR<T64, K> = self.0 .0.mod_switch::<Q2>();
let b: T64 = self.0 .1.mod_switch::<Q2>();
Self(GLWE(a, b))
}
}
// NOTE: the ugly const generics are temporary
pub fn blind_rotation<const N: usize, const K: usize, const KN: usize, const KN2: u64>(
c: TLWE<KN>,
btk: BootstrappingKey<N, K, KN>,
table: TGLWE<N, K>,
) -> TGLWE<N, K> {
let c_kn: TLWE<KN> = c.mod_switch::<KN2>();
let (a, b): (TR<T64, KN>, T64) = (c_kn.0 .0, c_kn.0 .1);
// two main parts: rotate by a known power of X, rotate by a secret
// power of X (using the C gate)
// table * X^-b, ie. left rotate
let v_xb: TGLWE<N, K> = table.left_rotate(b.0 as usize);
// rotate by a secret power of X using the cmux gate
let mut c_j: TGLWE<N, K> = v_xb.clone();
let _ = (1..K).map(|j| {
c_j = TGGSW::<N, K>::cmux(
btk.0[j].clone(),
c_j.clone(),
c_j.clone().left_rotate(a.0[j].0 as usize),
);
dbg!(&c_j);
});
c_j
}
#[derive(Clone, Debug)]
pub struct BootstrappingKey<const N: usize, const K: usize, const KN: usize>(pub Vec<TGGSW<N, K>>);
impl<const N: usize, const K: usize, const KN: usize> BootstrappingKey<N, K, KN> {
pub fn from_sk(mut rng: impl Rng, sk: &tglwe::SecretKey<N, K>) -> Result<Self> {
let (beta, l) = (2u32, 64u32); // TMP
//
let s: TR<Tn<N>, K> = sk.0 .0.clone();
// each btk_j = TGGSW_sk(s_i)
let btk: Vec<TGGSW<N, K>> = s
.iter()
.map(|s_i| TGGSW::<N, K>::encrypt_s(&mut rng, beta, l, sk, s_i))
.collect::<Result<Vec<_>>>()?;
Ok(Self(btk))
}
}
pub fn compute_lookup_table<const T: u64, const K: usize, const N: usize>() -> TGLWE<N, K> {
// from 2021-1402:
// v(x) = \sum_j^{N-1} [(p_j / 2N mod p)/p] X^j
// matrix of coefficients with size K*N = delta x T
let delta: usize = N / T as usize;
let values: Vec<Zq<T>> = (0..T).map(|v| Zq::<T>::from_u64(v)).collect();
let coeffs: Vec<Zq<T>> = (0..T as usize)
.flat_map(|i| vec![values[i]; delta])
.collect();
let table = Rq::<T, N>::from_vec(coeffs);
// encode the table as plaintext
let v: Tn<N> = TGLWE::<N, K>::encode::<T>(&table);
// encode the table as TGLWE ciphertext
let v: TGLWE<N, K> = TGLWE::<N, K>::from_plaintext(v);
v
}
impl<const K: usize> Add<TLWE<K>> for TLWE<K> {
@@ -170,6 +233,7 @@ impl<const K: usize> Mul<T64> for TLWE<K> {
mod tests {
use anyhow::Result;
use rand::distributions::Uniform;
use std::time::Instant;
use super::*;
@@ -186,9 +250,7 @@ mod tests {
let (sk, pk) = S::new_key(&mut rng)?;
let m = Rq::<T, 1>::rand_u64(&mut rng, msg_dist)?;
dbg!(&m);
let p: T64 = S::encode::<T>(&m);
dbg!(&p);
let c = S::encrypt(&mut rng, &pk, &p)?;
let p_recovered = c.decrypt(&sk);