arnaucube
arnaucube/fhe-studyPublic Self-hosted repository
tfhe: add blind_rotation & lookup table computation
arnaucube git@arnaucube.com
Sun, 03 Aug 2025 19:12:47 +0000
commit

2c20a2ed0e60afa5aac5713a26d2f0dbed069f08

parent

7bfcf6f7c118db228bd328487ec8d8e40261aacd

5 files changed, 101 insertions(+), 23 deletions(-)

jump to
M gfhe/src/glwe.rs → gfhe/src/glwe.rs

@@ -33,6 +33,9 @@ impl<R: Ring, const K: usize> GLWE<R, K> {

pub fn zero() -> Self { Self(TR::zero(), R::zero()) } + pub fn from_plaintext(p: R) -> Self { + Self(TR::zero(), p) + } pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<R, K>, PublicKey<R, K>)> { let Xi_key = Uniform::new(0_f64, 2_f64);
M tfhe/src/tglwe.rs → tfhe/src/tglwe.rs

@@ -19,12 +19,12 @@ 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; + 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 @@ self.0.decrypt(&sk.0)

} /// 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();

@@ -90,6 +93,11 @@ })

.collect::<Vec<T64>>(); 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))) } }

@@ -291,13 +299,14 @@ fn test_sample_extraction() -> Result<()> {

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 @@

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]); } }
M tfhe/src/tgsw.rs → tfhe/src/tgsw.rs

@@ -7,7 +7,10 @@

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]
M tfhe/src/tlev.rs → tfhe/src/tlev.rs

@@ -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)]
M tfhe/src/tlwe.rs → tfhe/src/tlwe.rs

@@ -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 @@ crate::tglwe::SecretKey(glwe::SecretKey::<Tn<N>, K>(TR(r)))

} } -// #[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 @@ }

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 @@ .sum();

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 @@ #[cfg(test)]

mod tests { use anyhow::Result; use rand::distributions::Uniform; + use std::time::Instant; use super::*;

@@ -186,9 +250,7 @@ for _ in 0..200 {

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);