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

polish & clean a bit

Browse Source
This commit is contained in:
arnaucube
2025-08-14 11:30:28 +00:00
parent bb3288f211
commit 12bb2aa3da
18 changed files with 98 additions and 261 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tfhe/src/lib.rs
View File

@@ -10,3 +10,5 @@ pub mod tglwe;
pub mod tgsw;
pub mod tlev;
pub mod tlwe;
pub(crate) const ERR_SIGMA: f64 = 3.2;

1
tfhe/src/tggsw.rs
View File

@@ -157,6 +157,7 @@ mod tests {
#[test]
fn test_external_product() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 64 },
k: 4,
t: 16, // plaintext modulus

22
tfhe/src/tglwe.rs
View File

@@ -1,22 +1,15 @@
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, RingParam, Rq, Tn, T64, TR};
use gfhe::{glwe, glwe::Param, 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(pub glwe::SecretKey<Tn>);
// pub struct SecretKey<const K: usize>(pub tlwe::SecretKey<K>);
impl SecretKey {
pub fn to_tlwe(&self, param: &Param) -> tlwe::SecretKey {
@@ -45,11 +38,9 @@ impl TGLWE {
}
pub fn new_key(mut rng: impl Rng, param: &Param) -> Result<(SecretKey, PublicKey)> {
// assert_eq!(KN, K * N); // this is wip, while not being able to compute K*N
let (sk_tlwe, _) = TLWE::new_key(&mut rng, &param.lwe())?; //param.lwe() so that it uses K*N
debug_assert_eq!(sk_tlwe.0 .0.r.len(), param.lwe().k); // =KN (sanity check)
// let sk = crate::tlwe::sk_to_tglwe::<N, K, KN>(sk_tlwe);
let sk = sk_tlwe.to_tglwe(param);
let pk: PublicKey = GLWE::pk_from_sk(rng, param, sk.0.clone())?;
Ok((sk, pk))
@@ -60,7 +51,6 @@ impl TGLWE {
let p = param.t;
let delta = u64::MAX / p; // floored
let coeffs = m.coeffs();
// Tn(array::from_fn(|i| T64(coeffs[i].0 * delta)))
Tn {
param: param.ring,
coeffs: coeffs.iter().map(|c_i| T64(c_i.v * delta)).collect(),
@@ -72,7 +62,7 @@ impl TGLWE {
Rq::from_vec_u64(&param.pt(), pt.coeffs().iter().map(|c| c.0).collect())
}
// encrypts with the given SecretKey (instead of PublicKey)
/// encrypts with the given SecretKey (instead of PublicKey)
pub fn encrypt_s(rng: impl Rng, param: &Param, sk: &SecretKey, p: &Tn) -> Result<Self> {
let glwe = GLWE::encrypt_s(rng, param, &sk.0, p)?;
Ok(Self(glwe))
@@ -141,11 +131,6 @@ impl Sum<TGLWE> for TGLWE {
where
I: Iterator<Item = Self>,
{
// let mut acc = TGLWE::<N, K>::zero();
// for e in iter {
// acc += e;
// }
// acc
let first = iter.next().unwrap();
iter.fold(first, |acc, e| acc + e)
}
@@ -208,6 +193,7 @@ mod tests {
#[test]
fn test_encrypt_decrypt() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 64 },
k: 16,
t: 128, // plaintext modulus
@@ -242,6 +228,7 @@ mod tests {
#[test]
fn test_addition() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 64 },
k: 16,
t: 128, // plaintext modulus
@@ -275,6 +262,7 @@ mod tests {
#[test]
fn test_add_plaintext() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 64 },
k: 16,
t: 128, // plaintext modulus
@@ -307,6 +295,7 @@ mod tests {
#[test]
fn test_mul_plaintext() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 64 },
k: 16,
t: 128, // plaintext modulus
@@ -342,6 +331,7 @@ mod tests {
#[test]
fn test_sample_extraction() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 64 },
k: 16,
t: 128, // plaintext modulus

17
tfhe/src/tgsw.rs
View File

@@ -1,17 +1,13 @@
use anyhow::Result;
use itertools::zip_eq;
use rand::Rng;
use std::array;
use std::ops::{Add, Mul};
use std::ops::Mul;
use arith::{Ring, RingParam, Rq, Tn, T64, TR};
use arith::{Ring, T64};
use crate::tlev::TLev;
use crate::{
tglwe::TGLWE,
tlwe::{PublicKey, SecretKey, TLWE},
};
use gfhe::glwe::{Param, GLWE};
use crate::tlwe::{SecretKey, TLWE};
use gfhe::glwe::Param;
/// vector of length K+1 = [K], [1]
#[derive(Clone, Debug)]
@@ -73,10 +69,12 @@ mod tests {
use rand::distributions::Uniform;
use super::*;
use arith::{RingParam, Rq};
#[test]
fn test_encrypt_decrypt() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 1 },
k: 16,
t: 2, // plaintext modulus
@@ -106,6 +104,7 @@ mod tests {
#[test]
fn test_external_product() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 1 },
k: 32,
t: 2, // plaintext modulus
@@ -130,7 +129,6 @@ mod tests {
let res: TLWE = tgsw * tlwe;
// let p_recovered = res.decrypt(&sk, beta);
let p_recovered = res.decrypt(&sk);
// downscaled by delta^-1
let res_recovered = TLWE::decode(&param, &p_recovered);
@@ -145,6 +143,7 @@ mod tests {
#[test]
fn test_cmux() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 1 },
k: 32,
t: 2, // plaintext modulus

30
tfhe/src/tlev.rs
View File

@@ -1,12 +1,10 @@
use anyhow::Result;
use itertools::zip_eq;
use rand::Rng;
use std::array;
use std::ops::{Add, Mul};
use std::ops::Mul;
use arith::{Ring, RingParam, Rq, Tn, T64, TR};
use arith::{Ring, RingParam, Rq, T64};
use crate::tglwe::TGLWE;
use crate::tlwe::{PublicKey, SecretKey, TLWE};
use gfhe::glwe::Param;
@@ -23,7 +21,6 @@ impl TLev {
}
pub fn decode(param: &Param, p: &T64) -> Rq {
Rq::from_vec_u64(
// &RingParam { q: u64::MAX, n: 1 },
&RingParam { q: param.t, n: 1 },
p.coeffs().iter().map(|c| c.0).collect(),
)
@@ -38,14 +35,16 @@ impl TLev {
) -> Result<Self> {
debug_assert_eq!(pk.1.k, param.k);
let tlev: Vec<TLWE> = (1..l + 1)
let tlev: Vec<TLWE> = (1..l as u64 + 1)
.map(|i| {
TLWE::encrypt(
&mut rng,
param,
pk,
&(*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::encrypt(&mut rng, param, pk, &aux)
})
.collect::<Result<Vec<_>>>()?;
@@ -59,6 +58,8 @@ impl TLev {
sk: &SecretKey,
m: &T64,
) -> Result<Self> {
debug_assert_eq!(sk.0 .0.k, param.k);
let tlev: Vec<TLWE> = (1..l as u64 + 1)
.map(|i| {
let aux = if i < 64 {
@@ -113,6 +114,7 @@ mod tests {
#[test]
fn test_encrypt_decrypt() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 1 },
k: 16,
t: 2, // plaintext modulus
@@ -143,13 +145,15 @@ mod tests {
#[test]
fn test_tlev_vect64_product() -> Result<()> {
let param = Param {
err_sigma: 0.1, // WIP
ring: RingParam { q: u64::MAX, n: 1 },
k: 16,
t: 2, // plaintext modulus
};
let beta: u32 = 2;
let l: u32 = 16;
// let l: u32 = 16;
let l: u32 = 64;
let mut rng = rand::thread_rng();
let msg_dist = Uniform::new(0_u64, param.t);

29
tfhe/src/tlwe.rs
View File

@@ -49,11 +49,10 @@ impl TLWE {
Ok((SecretKey(sk), pk))
}
// TODO use param instead of p:u64
pub fn encode(param: &Param, m: &Rq) -> T64 {
assert_eq!(param.ring.n, 1);
debug_assert_eq!(param.t, m.param.q); // plaintext modulus
//
let delta = u64::MAX / param.t; // floored
let coeffs = m.coeffs();
T64(coeffs[0].v * delta)
@@ -112,14 +111,14 @@ impl TLWE {
Self(GLWE(a, b))
}
}
// NOTE: the ugly const generics are temporary
pub fn blind_rotation(
param: &Param,
c: TLWE, // kn
btk: BootstrappingKey,
table: TGLWE, // n,k
) -> TGLWE {
debug_assert_eq!(c.0 .0.k, param.k); // TODO or k*n?
debug_assert_eq!(c.0 .0.k, param.k);
// TODO replace `param.k*param.ring.n` by `param.kn()`
let c_kn: TLWE = c.mod_switch((param.k * param.ring.n) as u64);
@@ -138,7 +137,6 @@ pub fn blind_rotation(
c_j.clone(),
c_j.clone().left_rotate(a.r[j].0 as usize),
);
// dbg!(&c_j);
});
c_j
}
@@ -164,7 +162,7 @@ pub struct BootstrappingKey(
impl BootstrappingKey {
pub fn from_sk(mut rng: impl Rng, param: &Param, sk: &tglwe::SecretKey) -> Result<Self> {
let (beta, l) = (2u32, 64u32); // TMP
//
let s: TR<Tn> = sk.0 .0.clone();
let (sk2, _) = TLWE::new_key(&mut rng, &param.lwe())?; // TLWE<KN> compatible with TGLWE<N,K>
@@ -229,11 +227,6 @@ impl Sum<TLWE> for TLWE {
where
I: Iterator<Item = Self>,
{
// let mut acc = TLWE::<K>::zero();
// for e in iter {
// acc += e;
// }
// acc
let first = iter.next().unwrap();
iter.fold(first, |acc, e| acc + e)
}
@@ -290,6 +283,7 @@ mod tests {
#[test]
fn test_encrypt_decrypt() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 1 },
k: 16,
t: 128, // plaintext modulus
@@ -324,6 +318,7 @@ mod tests {
#[test]
fn test_addition() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 1 },
k: 16,
t: 128, // plaintext modulus
@@ -357,6 +352,7 @@ mod tests {
#[test]
fn test_add_plaintext() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 1 },
k: 16,
t: 128, // plaintext modulus
@@ -389,6 +385,7 @@ mod tests {
#[test]
fn test_mul_plaintext() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 1 },
k: 16,
t: 128, // plaintext modulus
@@ -404,7 +401,6 @@ mod tests {
let m2 = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
let p1: T64 = TLWE::encode(&param, &m1);
// don't scale up p2, set it directly from m2
// let p2: T64 = Tn(array::from_fn(|i| T64(m2.coeffs()[i].0)));
let p2: T64 = T64(m2.coeffs()[0].v);
let c1 = TLWE::encrypt(&mut rng, &param, &pk, &p1)?;
@@ -422,6 +418,7 @@ mod tests {
#[test]
fn test_key_switch() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam { q: u64::MAX, n: 1 },
k: 16,
t: 128, // plaintext modulus
@@ -440,7 +437,7 @@ mod tests {
let msg_dist = Uniform::new(0_u64, param.t);
let m = Rq::rand_u64(&mut rng, msg_dist, &param.pt())?;
let p = TLWE::encode(&param, &m); // plaintext
//
let c = TLWE::encrypt_s(&mut rng, &param, &sk, &p)?;
let c2 = c.key_switch(&param, beta, l, &ksk);
@@ -463,6 +460,7 @@ mod tests {
#[test]
fn test_bootstrapping() -> Result<()> {
let param = Param {
err_sigma: crate::ERR_SIGMA,
ring: RingParam {
q: u64::MAX,
n: 1024,
@@ -470,10 +468,6 @@ mod tests {
k: 1,
t: 128, // plaintext modulus
};
// const T: u64 = 128; // plaintext modulus
// const K: usize = 1;
// const N: usize = 1024;
// const KN: usize = K * N;
let mut rng = rand::thread_rng();
let start = Instant::now();
@@ -494,7 +488,6 @@ mod tests {
let c = TLWE::encrypt_s(&mut rng, &param.lwe(), &sk_tlwe, &p)?;
let start = Instant::now();
// the ugly const generics are temporary
let bootstrapped: TLWE = bootstrapping(&param, btk, table, c);
println!("bootstrapping took: {:?}", start.elapsed());