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

ckks: get rid of constant generics (reason in two commits ago)

Browse Source
This commit is contained in:
arnaucube
2025-08-11 21:40:58 +00:00
parent 78324a3a8d
commit 0c7e078aeb
3 changed files with 113 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files

29
ckks/src/encoder.rs
View File

@@ -3,12 +3,13 @@ use anyhow::Result;
use arith::{Matrix, Ring, Rq, C, R};
#[derive(Clone, Debug)]
pub struct SecretKey<const Q: u64, const N: usize>(Rq<Q, N>);
pub struct SecretKey(Rq);
#[derive(Clone, Debug)]
pub struct PublicKey<const Q: u64, const N: usize>(Rq<Q, N>, Rq<Q, N>);
pub struct PublicKey(Rq, Rq);
pub struct Encoder<const Q: u64, const N: usize> {
pub struct Encoder {
n: usize,
scale_factor: C<f64>, // Δ (delta)
primitive: C<f64>,
basis: Matrix<C<f64>>,
@@ -34,13 +35,14 @@ fn vandermonde(n: usize, w: C<f64>) -> Matrix<C<f64>> {
}
Matrix::<C<f64>>(v)
}
impl<const Q: u64, const N: usize> Encoder<Q, N> {
pub fn new(scale_factor: C<f64>) -> Self {
let primitive: C<f64> = primitive_root_of_unity(2 * N);
let basis = vandermonde(N, primitive);
impl Encoder {
pub fn new(n: usize, scale_factor: C<f64>) -> Self {
let primitive: C<f64> = primitive_root_of_unity(2 * n);
let basis = vandermonde(n, primitive);
let basis_t = basis.transpose();
Self {
n,
scale_factor,
primitive,
basis,
@@ -52,7 +54,7 @@ impl<const Q: u64, const N: usize> Encoder<Q, N> {
/// from $\mathbb{C}^{N/2} \longrightarrow \mathbb{Z_q}[X]/(X^N +1) = R$
// TODO use alg.1 from 2018-1043,
// or as in 2018-1073: $f(x) = 1N (U^T.conj() m + U^T m.conj())$
pub fn encode(&self, z: &[C<f64>]) -> Result<R<N>> {
pub fn encode(&self, z: &[C<f64>]) -> Result<R> {
// $pi^{-1}: \mathbb{C}^{N/2} \longrightarrow \mathbb{H}$
let expanded = self.pi_inv(z);
@@ -93,10 +95,10 @@ impl<const Q: u64, const N: usize> Encoder<Q, N> {
// TMP: naive round, maybe do gaussian
let coeffs = r.iter().map(|e| e.re.round() as i64).collect::<Vec<i64>>();
Ok(R::from_vec(coeffs))
Ok(R::from_vec(self.n, coeffs))
}
pub fn decode(&self, p: &R<N>) -> Result<Vec<C<f64>>> {
pub fn decode(&self, p: &R) -> Result<Vec<C<f64>>> {
let p: Vec<C<f64>> = p
.coeffs()
.iter()
@@ -110,7 +112,7 @@ impl<const Q: u64, const N: usize> Encoder<Q, N> {
/// pi: \mathbb{H} \longrightarrow \mathbb{C}^{N/2}
fn pi(&self, z: &[C<f64>]) -> Vec<C<f64>> {
z[..N / 2].to_vec()
z[..self.n / 2].to_vec()
}
/// pi^{-1}: \mathbb{C}^{N/2} \longrightarrow \mathbb{H}
fn pi_inv(&self, z: &[C<f64>]) -> Vec<C<f64>> {
@@ -154,6 +156,7 @@ mod tests {
fn test_encode_decode() -> Result<()> {
const Q: u64 = 1024;
const N: usize = 32;
let n: usize = 32;
let T = 128; // WIP
let mut rng = rand::thread_rng();
@@ -166,9 +169,9 @@ mod tests {
.collect();
let delta = C::<f64>::new(64.0, 0.0); // delta = scaling factor
let encoder = Encoder::<Q, N>::new(delta);
let encoder = Encoder::new(n, delta);
let m: R<N> = encoder.encode(&z)?; // polynomial (encoded vec) \in R
let m: R = encoder.encode(&z)?; // polynomial (encoded vec) \in R
let z_decoded = encoder.decode(&m)?;

170
ckks/src/lib.rs
View File

@@ -18,35 +18,48 @@ pub use encoder::Encoder;
// sigma=3.2 from: https://eprint.iacr.org/2016/421.pdf page 17
const ERR_SIGMA: f64 = 3.2;
#[derive(Debug)]
pub struct PublicKey<const Q: u64, const N: usize>(Rq<Q, N>, Rq<Q, N>);
pub struct SecretKey<const Q: u64, const N: usize>(Rq<Q, N>);
pub struct CKKS<const Q: u64, const N: usize> {
encoder: Encoder<Q, N>,
#[derive(Clone, Copy, Debug)]
pub struct Params {
q: u64,
n: usize,
t: u64,
}
impl<const Q: u64, const N: usize> CKKS<Q, N> {
pub fn new(delta: C<f64>) -> Self {
let encoder = Encoder::<Q, N>::new(delta);
Self { encoder }
#[derive(Debug)]
pub struct PublicKey(Rq, Rq);
pub struct SecretKey(Rq);
pub struct CKKS {
params: Params,
encoder: Encoder,
}
impl CKKS {
pub fn new(params: &Params, delta: C<f64>) -> Self {
let encoder = Encoder::new(params.n, delta);
Self {
params: params.clone(),
encoder,
}
}
/// generate a new key pair (privK, pubK)
pub fn new_key(&self, mut rng: impl Rng) -> Result<(SecretKey<Q, N>, PublicKey<Q, N>)> {
pub fn new_key(&self, mut rng: impl Rng) -> Result<(SecretKey, PublicKey)> {
let params = &self.params;
let Xi_key = Uniform::new(-1_f64, 1_f64);
let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
let e = Rq::<Q, N>::rand_f64(&mut rng, Xi_err)?;
let e = Rq::rand_f64(&mut rng, Xi_err, params.q, params.n)?;
let mut s = Rq::<Q, N>::rand_f64(&mut rng, Xi_key)?;
let mut s = Rq::rand_f64(&mut rng, Xi_key, params.q, params.n)?;
// since s is going to be multiplied by other Rq elements, already
// compute its NTT
s.compute_evals();
let a = Rq::<Q, N>::rand_f64(&mut rng, Xi_key)?;
let a = Rq::rand_f64(&mut rng, Xi_key, params.q, params.n)?;
let pk: PublicKey<Q, N> = PublicKey((&(-a) * &s) + e, a.clone());
let pk: PublicKey = PublicKey((&(-a.clone()) * &s) + e, a.clone()); // TODO rm clones
Ok((SecretKey(s), pk))
}
@@ -54,64 +67,54 @@ impl<const Q: u64, const N: usize> CKKS<Q, N> {
fn encrypt(
&self, // TODO maybe rm?
mut rng: impl Rng,
pk: &PublicKey<Q, N>,
m: &R<N>,
) -> Result<(Rq<Q, N>, Rq<Q, N>)> {
pk: &PublicKey,
m: &R,
) -> Result<(Rq, Rq)> {
let params = self.params;
let Xi_key = Uniform::new(-1_f64, 1_f64);
let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
let e_0 = Rq::<Q, N>::rand_f64(&mut rng, Xi_err)?;
let e_1 = Rq::<Q, N>::rand_f64(&mut rng, Xi_err)?;
let e_0 = Rq::rand_f64(&mut rng, Xi_err, params.q, params.n)?;
let e_1 = Rq::rand_f64(&mut rng, Xi_err, params.q, params.n)?;
let v = Rq::<Q, N>::rand_f64(&mut rng, Xi_key)?;
let v = Rq::rand_f64(&mut rng, Xi_key, params.q, params.n)?;
let m: Rq<Q, N> = Rq::<Q, N>::from(*m);
// let m: Rq = Rq::from(*m);
let m: Rq = m.clone().to_rq(params.q); // TODO rm clone
Ok((m + e_0 + v * pk.0.clone(), v * pk.1.clone() + e_1))
Ok((m + e_0 + &v * &pk.0.clone(), &v * &pk.1 + e_1))
}
fn decrypt(
&self, // TODO maybe rm?
sk: &SecretKey<Q, N>,
c: (Rq<Q, N>, Rq<Q, N>),
) -> Result<R<N>> {
let m = c.0.clone() + c.1 * sk.0;
sk: &SecretKey,
c: (Rq, Rq),
) -> Result<R> {
let m = c.0.clone() + &c.1 * &sk.0;
Ok(m.mod_centered_q())
}
pub fn encode_and_encrypt(
&self,
mut rng: impl Rng,
pk: &PublicKey<Q, N>,
pk: &PublicKey,
z: &[C<f64>],
) -> Result<(Rq<Q, N>, Rq<Q, N>)> {
let m: R<N> = self.encoder.encode(&z)?; // polynomial (encoded vec) \in R
) -> Result<(Rq, Rq)> {
let m: R = self.encoder.encode(&z)?; // polynomial (encoded vec) \in R
self.encrypt(&mut rng, pk, &m)
}
pub fn decrypt_and_decode(
&self,
sk: SecretKey<Q, N>,
c: (Rq<Q, N>, Rq<Q, N>),
) -> Result<Vec<C<f64>>> {
pub fn decrypt_and_decode(&self, sk: SecretKey, c: (Rq, Rq)) -> Result<Vec<C<f64>>> {
let d = self.decrypt(&sk, c)?;
self.encoder.decode(&d)
}
pub fn add(
&self,
c0: &(Rq<Q, N>, Rq<Q, N>),
c1: &(Rq<Q, N>, Rq<Q, N>),
) -> Result<(Rq<Q, N>, Rq<Q, N>)> {
pub fn add(&self, c0: &(Rq, Rq), c1: &(Rq, Rq)) -> Result<(Rq, Rq)> {
Ok((&c0.0 + &c1.0, &c0.1 + &c1.1))
}
pub fn sub(
&self,
c0: &(Rq<Q, N>, Rq<Q, N>),
c1: &(Rq<Q, N>, Rq<Q, N>),
) -> Result<(Rq<Q, N>, Rq<Q, N>)> {
pub fn sub(&self, c0: &(Rq, Rq), c1: &(Rq, Rq)) -> Result<(Rq, Rq)> {
Ok((&c0.0 - &c1.0, &c0.1 + &c1.1))
}
}
@@ -122,21 +125,22 @@ mod tests {
#[test]
fn test_encrypt_decrypt() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 32;
const T: u64 = 50;
let q: u64 = 2u64.pow(16) + 1;
let n: usize = 32;
let t: u64 = 50;
let params = Params { q, n, t };
let scale_factor_u64 = 512_u64; // delta
let scale_factor = C::<f64>::new(scale_factor_u64 as f64, 0.0); // delta
let mut rng = rand::thread_rng();
for _ in 0..1000 {
let ckks = CKKS::<Q, N>::new(scale_factor);
let ckks = CKKS::new(&params, scale_factor);
let (sk, pk) = ckks.new_key(&mut rng)?;
let m_raw: R<N> = Rq::<Q, N>::rand_f64(&mut rng, Uniform::new(0_f64, T as f64))?.to_r();
let m = m_raw * scale_factor_u64;
let m_raw: R = Rq::rand_f64(&mut rng, Uniform::new(0_f64, t as f64), q, n)?.to_r();
let m = &m_raw * &scale_factor_u64;
let ct = ckks.encrypt(&mut rng, &pk, &m)?;
let m_decrypted = ckks.decrypt(&sk, ct)?;
@@ -146,8 +150,9 @@ mod tests {
.iter()
.map(|e| (*e as f64 / (scale_factor_u64 as f64)).round() as u64)
.collect();
let m_decrypted = Rq::<Q, N>::from_vec_u64(m_decrypted);
assert_eq!(m_decrypted, Rq::<Q, N>::from(m_raw));
let m_decrypted = Rq::from_vec_u64(q, n, m_decrypted);
// assert_eq!(m_decrypted, Rq::from(m_raw));
assert_eq!(m_decrypted, m_raw.to_rq(q));
}
Ok(())
@@ -155,21 +160,22 @@ mod tests {
#[test]
fn test_encode_encrypt_decrypt_decode() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 16;
const T: u64 = 8;
let q: u64 = 2u64.pow(16) + 1;
let n: usize = 16;
let t: u64 = 8;
let params = Params { q, n, t };
let scale_factor = C::<f64>::new(512.0, 0.0); // delta
let mut rng = rand::thread_rng();
for _ in 0..1000 {
let ckks = CKKS::<Q, N>::new(scale_factor);
let ckks = CKKS::new(&params, scale_factor);
let (sk, pk) = ckks.new_key(&mut rng)?;
let z: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, T))
.take(N / 2)
let z: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, t))
.take(n / 2)
.collect();
let m: R<N> = ckks.encoder.encode(&z)?;
let m: R = ckks.encoder.encode(&z)?;
println!("{}", m);
// sanity check
@@ -200,26 +206,27 @@ mod tests {
#[test]
fn test_add() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 16;
const T: u64 = 8;
let q: u64 = 2u64.pow(16) + 1;
let n: usize = 16;
let t: u64 = 8;
let params = Params { q, n, t };
let scale_factor = C::<f64>::new(1024.0, 0.0); // delta
let mut rng = rand::thread_rng();
for _ in 0..1000 {
let ckks = CKKS::<Q, N>::new(scale_factor);
let ckks = CKKS::new(&params, scale_factor);
let (sk, pk) = ckks.new_key(&mut rng)?;
let z0: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, T))
.take(N / 2)
let z0: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, t))
.take(n / 2)
.collect();
let z1: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, T))
.take(N / 2)
let z1: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, t))
.take(n / 2)
.collect();
let m0: R<N> = ckks.encoder.encode(&z0)?;
let m1: R<N> = ckks.encoder.encode(&z1)?;
let m0: R = ckks.encoder.encode(&z0)?;
let m1: R = ckks.encoder.encode(&z1)?;
let ct0 = ckks.encrypt(&mut rng, &pk, &m0)?;
let ct1 = ckks.encrypt(&mut rng, &pk, &m1)?;
@@ -243,26 +250,27 @@ mod tests {
#[test]
fn test_sub() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 16;
const T: u64 = 8;
let q: u64 = 2u64.pow(16) + 1;
let n: usize = 16;
let t: u64 = 8;
let params = Params { q, n, t };
let scale_factor = C::<f64>::new(1024.0, 0.0); // delta
let mut rng = rand::thread_rng();
for _ in 0..1000 {
let ckks = CKKS::<Q, N>::new(scale_factor);
let ckks = CKKS::new(&params, scale_factor);
let (sk, pk) = ckks.new_key(&mut rng)?;
let z0: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, T))
.take(N / 2)
let z0: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, t))
.take(n / 2)
.collect();
let z1: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, T))
.take(N / 2)
let z1: Vec<C<f64>> = std::iter::repeat_with(|| C::<f64>::rand(&mut rng, t))
.take(n / 2)
.collect();
let m0: R<N> = ckks.encoder.encode(&z0)?;
let m1: R<N> = ckks.encoder.encode(&z1)?;
let m0: R = ckks.encoder.encode(&z0)?;
let m1: R = ckks.encoder.encode(&z1)?;
let ct0 = ckks.encrypt(&mut rng, &pk, &m0)?;
let ct1 = ckks.encrypt(&mut rng, &pk, &m1)?;