arnaucube
arnaucube/fhe-studyPublic Self-hosted repository
implement GLWE key switching
arnaucube git@arnaucube.com
Wed, 16 Jul 2025 21:02:07 +0200
commit

188bc7fa7f13ebd88f53b5f4783c417e5c0c8d6c

parent

1e2ea824fd50d21fe2c52532e422821508415317

A arith/sage/ring.sage

@@ -0,0 +1,77 @@

+def run_test(a, b): + print("\nnew test:") + print(a) + print(b) + c = a*b + print(c) + print(c.list()) + + + +n_iters = 100 +Q= 65537 +print(Q) +N=4 +F = GF(Q) +R = QuotientRing(F[x], x^N + 1, names="X") +print(R) + +a = R([4,2,1,0]) +b = R([1,2,3,4]) +run_test(a,b) + + +# print("Elements of the polynomial ring:") +# for e in R: +# print(e) + + +# Other: +# ====== +# +# t = R.gen() +# a = 0 + t + 2*t^2 + 3*t^3 + 4*t^4 + 5*t^5 +# b = 5 + 4*t + 3*t^2 + 2*t^3 + 1*t^4 + 0*t^5 +# print("add", a+b) +# print("sub", a-b) +# print("mul", a*b) +# a = 0 + t + 2*t^2 + 3*t^3 + 4*t^4 + 5*t^5 + +# print("ring elem mul testvectors") +# +# +# def randvec(size=N):return [int(random()*(Q-1)) for t in range(size)] +# +# a_vecs = [None]*n_iters +# b_vecs = [None]*n_iters +# c_vecs = [None]*n_iters +# +# for i in range(n_iters): +# a_vec = randvec() +# b_vec = randvec() +# a_pol = R(a_vec) +# b_pol = R(b_vec) +# +# c_pol = a_pol*b_pol +# +# a_vecs[i] = a_pol.list() +# b_vecs[i] = b_pol.list() +# c_vecs[i] = c_pol.list() +# +# print("let a_vecs = vec!{};\n".format(a_vecs)) +# print("let b_vecs = vec!{};\n".format(b_vecs)) +# print("let c_vecs = vec!{};".format(c_vecs)) + + +# # cyclotomic +# +# Q= 65537 +# print(Q) +# N=4 +# F = GF(Q) +# R = QuotientRing(F[x], x^N - 1, names="X") +# print(R) +# +# a = R([1,0,0,2]) +# b = R([0,0,0,2]) +# run_test(a, b)
M arith/src/ring.rs → arith/src/ring.rs

@@ -31,6 +31,12 @@ Self(coeffs)

// let coeffs: [C; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng))); // Self(coeffs) } + + // returns the decomposition of each polynomial coefficient + fn decompose(&self, beta: u32, l: u32) -> Vec<Self> { + unimplemented!(); + // array::from_fn(|i| self.coeffs[i].decompose(beta, l)) + } } impl<const Q: u64, const N: usize> From<crate::ringq::Rq<Q, N>> for R<N> {
M arith/src/ringq.rs → arith/src/ringq.rs

@@ -47,6 +47,18 @@ coeffs,

evals: None, } } + + // returns the decomposition of each polynomial coefficient, such + // decomposition will be a vecotor of length N, containint N vectors of Zq + fn decompose(&self, beta: u32, l: u32) -> Vec<Self> { + let elems: Vec<Vec<Zq<Q>>> = self.coeffs.iter().map(|r| r.decompose(beta, l)).collect(); + // transpose it + let r: Vec<Vec<Zq<Q>>> = (0..elems[0].len()) + .map(|i| (0..elems.len()).map(|j| elems[j][i]).collect()) + .collect(); + // convert it to Rq<Q,N> + r.iter().map(|a_i| Self::from_vec(a_i.clone())).collect() + } } impl<const Q: u64, const N: usize> From<crate::ring::R<N>> for Rq<Q, N> {

@@ -597,6 +609,33 @@ let expected_c = Rq::new(expected_c, None);

let c = mul_mut(&mut a, &mut b); assert_eq!(c, expected_c); + Ok(()) + } + + #[test] + fn test_rq_decompose() -> Result<()> { + const Q: u64 = 16; + const N: usize = 4; + let beta = 4; + let l = 2; + + let a = Rq::<Q, N>::from_vec_u64(vec![7u64, 14, 3, 6]); + let d = a.decompose(beta, l); + + assert_eq!( + d[0], + vec![1u64, 3, 0, 1] + .iter() + .map(|e| Zq::<Q>::from_u64(*e)) + .collect::<Vec<_>>() + ); + assert_eq!( + d[1], + vec![3u64, 2, 3, 2] + .iter() + .map(|e| Zq::<Q>::from_u64(*e)) + .collect::<Vec<_>>() + ); Ok(()) } }
M arith/src/traits.rs → arith/src/traits.rs

@@ -27,4 +27,6 @@ fn coeffs(&self) -> Vec<Self::C>;

fn zero() -> Self; // note/wip/warning: dist (0,q) with f64, will output more '0=q' elements than other values fn rand(rng: impl Rng, dist: impl Distribution<f64>) -> Self; + + fn decompose(&self, beta: u32, l: u32) -> Vec<Self>; }
M arith/src/tuple_ring.rs → arith/src/tuple_ring.rs

@@ -16,6 +16,9 @@ #[derive(Clone, Debug)]

pub struct TR<R: Ring, const K: usize>(pub Vec<R>); impl<R: Ring, const K: usize> TR<R, K> { + pub fn zero() -> Self { + Self((0..K).into_iter().map(|_| R::zero()).collect()) + } pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self { Self( (0..K)

@@ -23,6 +26,11 @@ .into_iter()

.map(|_| R::rand(&mut rng, &dist)) .collect(), ) + } + // returns the decomposition of each polynomial element + pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> { + unimplemented!() + // self.0.iter().map(|r| r.decompose(beta, l)).collect() // this is Vec<Vec<Vec<R::C>>> } }
M arith/src/zq.rs → arith/src/zq.rs

@@ -3,7 +3,7 @@ use rand::{distributions::Distribution, Rng};

use std::fmt; use std::ops; -// Z_q, integers modulus q, not necessarily prime +/// Z_q, integers modulus q, not necessarily prime #[derive(Clone, Copy, PartialEq)] pub struct Zq<const Q: u64>(pub u64);

@@ -130,6 +130,28 @@ /// perform the mod switch operation from Q to Q', where Q2=Q'

pub fn mod_switch<const Q2: u64>(&self) -> Zq<Q2> { Zq::<Q2>::from_u64(((self.0 as f64 * Q2 as f64) / Q as f64).round() as u64) } + + // TODO more efficient method for when decomposing with base 2 (beta=2) + pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> { + let mut rem: u64 = self.0; + // next if is for cases in which beta does not divide Q (concretely + // beta^l!=Q). round to the nearest multiple of q/beta^l + if rem >= beta.pow(l) as u64 { + // rem = Q - 1 - (Q / beta as u64); // floor + return vec![Zq(beta as u64 - 1); l as usize]; + } + + let mut x: Vec<Self> = vec![]; + for i in 1..l + 1 { + let den = Q / beta.pow(i) as u64; + let x_i = rem / den; // division between u64 already does floor + x.push(Self::from_u64(x_i)); + if x_i != 0 { + rem = rem % den; + } + } + x + } } impl<const Q: u64> Zq<Q> {

@@ -243,6 +265,7 @@

#[cfg(test)] mod tests { use super::*; + use rand::distributions::Uniform; #[test] fn exp() {

@@ -261,5 +284,63 @@ const Q: u64 = 1021;

let a = Zq::<Q>::from_f64(101.0); let b = Zq::<Q>::from_f64(-1.0); assert_eq!(-a, a * b); + } + + fn recompose<const Q: u64>(beta: u32, l: u32, d: Vec<Zq<Q>>) -> Zq<Q> { + let mut x = 0u64; + for i in 0..l { + x += d[i as usize].0 * Q / beta.pow(i + 1) as u64; + } + Zq::from_u64(x) + } + + #[test] + fn test_decompose() { + const Q1: u64 = 16; + let beta: u32 = 2; + let l: u32 = 4; + let x = Zq::<Q1>::from_u64(9); + let d = x.decompose(beta, l); + + assert_eq!(recompose::<Q1>(beta, l, d), x); + + const Q: u64 = 5u64.pow(3); + let beta: u32 = 5; + let l: u32 = 3; + + let dist = Uniform::new(0_u64, Q); + let mut rng = rand::thread_rng(); + + for _ in 0..1000 { + let x = Zq::<Q>::from_u64(dist.sample(&mut rng)); + + let d = x.decompose(beta, l); + + assert_eq!(recompose::<Q>(beta, l, d), x); + } + } + + #[test] + fn test_decompose_approx() { + const Q: u64 = 2u64.pow(4) + 1; + let beta: u32 = 2; + let l: u32 = 4; + let x = Zq::<Q>::from_u64(16); // in q, but bigger than beta^l + let d = x.decompose(beta, l); + assert_eq!(recompose::<Q>(beta, l, d), Zq(15)); + + const Q2: u64 = 5u64.pow(3) + 1; + let beta: u32 = 5; + let l: u32 = 3; + let x = Zq::<Q2>::from_u64(125); // in q, but bigger than beta^l + let d = x.decompose(beta, l); + assert_eq!(recompose::<Q2>(beta, l, d), Zq(124)); + + const Q3: u64 = 2u64.pow(16) + 1; + let beta: u32 = 2; + let l: u32 = 16; + let x = Zq::<Q3>::from_u64(Q3 - 1); // in q, but bigger than beta^l + let d = x.decompose(beta, l); + assert_eq!(recompose::<Q3>(beta, l, d), Zq(beta.pow(l) as u64 - 1)); } }
M gfhe/src/glwe.rs → gfhe/src/glwe.rs

@@ -1,20 +1,33 @@

use anyhow::Result; +use itertools::zip_eq; use rand::Rng; use rand_distr::{Normal, Uniform}; -use std::ops::{Add, Mul}; +use std::iter::Sum; +use std::ops::{Add, AddAssign, Mul, Sub}; + +use arith::{Ring, Rq, Zq, TR}; -use arith::{Ring, Rq, TR}; +use crate::glev::GLev; const ERR_SIGMA: f64 = 3.2; +#[derive(Clone, Debug)] pub struct GLWE<const Q: u64, const N: usize, const K: usize>(TR<Rq<Q, N>, K>, Rq<Q, N>); #[derive(Clone, Debug)] pub struct SecretKey<const Q: u64, const N: usize, const K: usize>(TR<Rq<Q, N>, K>); #[derive(Clone, Debug)] pub struct PublicKey<const Q: u64, const N: usize, const K: usize>(Rq<Q, N>, TR<Rq<Q, N>, K>); + +// K GLevs, each KSK_i=l GLWEs +#[derive(Clone, Debug)] +pub struct KSK<const Q: u64, const N: usize, const K: usize>(Vec<GLev<Q, N, K>>); impl<const Q: u64, const N: usize, const K: usize> GLWE<Q, N, K> { + pub fn zero() -> Self { + Self(TR::zero(), Rq::zero()) + } + pub fn new_key(mut rng: impl Rng) -> Result<(SecretKey<Q, N, K>, PublicKey<Q, N, K>)> { let Xi_key = Uniform::new(0_f64, 2_f64); let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;

@@ -27,32 +40,69 @@ let pk: PublicKey<Q, N, K> = PublicKey((&a * &s) + e, a);

Ok((SecretKey(s), pk)) } - // TODO delta not as input - pub fn encrypt_s<const T: u64>( + pub fn new_ksk( + mut rng: impl Rng, + beta: u32, + l: u32, + sk: &SecretKey<Q, N, K>, + new_sk: &SecretKey<Q, N, K>, + ) -> Result<KSK<Q, N, K>> { + let r: Vec<GLev<Q, N, K>> = (0..K) + .into_iter() + .map(|i| + // treat sk_i as the msg being encrypted + GLev::<Q, N, K>::encrypt_s(&mut rng, beta, l, &new_sk, &sk.0 .0[i])) + .collect::<Result<Vec<_>>>()?; + + Ok(KSK(r)) + } + pub fn key_switch(&self, beta: u32, l: u32, ksk: &KSK<Q, N, K>) -> Self { + let (a, b): (TR<Rq<Q, N>, K>, Rq<Q, N>) = (self.0.clone(), self.1); + + let lhs: GLWE<Q, N, K> = GLWE(TR::zero(), b); + + // K iterations, ksk.0 contains K times GLev + let rhs: GLWE<Q, N, K> = zip_eq(a.0, ksk.0.clone()) + .map(|(a_i, ksk_i)| Self::dot_prod(a_i.decompose(beta, l), ksk_i)) + .sum(); + + lhs - rhs + } + // note: a_decomp is of length N + fn dot_prod(a_decomp: Vec<Rq<Q, N>>, ksk_i: GLev<Q, N, K>) -> GLWE<Q, N, K> { + // l times GLWES + let glwes: Vec<GLWE<Q, N, K>> = ksk_i.0; + + // l iterations + let r: GLWE<Q, N, K> = zip_eq(a_decomp, glwes) + .map(|(a_d_i, glwe_i)| glwe_i * a_d_i) + .sum(); + r + } + + // encrypts with the given SecretKey (instead of PublicKey) + pub fn encrypt_s( mut rng: impl Rng, sk: &SecretKey<Q, N, K>, - m: &Rq<T, N>, + m: &Rq<Q, N>, + // TODO delta not as input delta: u64, ) -> Result<Self> { - let m: Rq<Q, N> = m.remodule::<Q>(); - let Xi_key = Uniform::new(0_f64, 2_f64); let Xi_err = Normal::new(0_f64, ERR_SIGMA)?; let a: TR<Rq<Q, N>, K> = TR::rand(&mut rng, Xi_key); let e = Rq::<Q, N>::rand(&mut rng, Xi_err); - let b: Rq<Q, N> = (&a * &sk.0) + m * delta + e; + let b: Rq<Q, N> = (&a * &sk.0) + *m * delta + e; Ok(Self(a, b)) } - pub fn encrypt<const T: u64>( + pub fn encrypt( mut rng: impl Rng, pk: &PublicKey<Q, N, K>, - m: &Rq<T, N>, + m: &Rq<Q, N>, delta: u64, ) -> Result<Self> { - let m: Rq<Q, N> = m.remodule::<Q>(); - let Xi_key = Uniform::new(0_f64, 2_f64); let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;

@@ -61,23 +111,16 @@

let e0 = Rq::<Q, N>::rand(&mut rng, Xi_err); let e1 = TR::<Rq<Q, N>, K>::rand(&mut rng, Xi_err); - let b: Rq<Q, N> = pk.0 * u + m * delta + e0; + let b: Rq<Q, N> = pk.0 * u + *m * delta + e0; let d: TR<Rq<Q, N>, K> = &pk.1 * &u + e1; Ok(Self(d, b)) } - pub fn decrypt<const T: u64>(&self, sk: &SecretKey<Q, N, K>, delta: u64) -> Rq<T, N> { + pub fn decrypt<const T: u64>(&self, sk: &SecretKey<Q, N, K>, delta: u64) -> Rq<Q, N> { let (d, b): (TR<Rq<Q, N>, K>, Rq<Q, N>) = (self.0.clone(), self.1); let r: Rq<Q, N> = b - &d * &sk.0; let r = r.mul_div_round(T, Q); - // let r_scaled: Vec<f64> = r - // .coeffs() - // .iter() - // // .map(|e| (e.0 as f64 / delta as f64).round()) - // .map(|e| e.mul_div_round(T, Q)) - // .collect(); - // let r = Rq::<Q, N>::from_vec_f64(r_scaled); - r.remodule::<T>() + r } pub fn mod_switch<const P: u64>(&self) -> GLWE<P, N, K> {

@@ -104,6 +147,36 @@ let b: Rq<Q, N> = self.1 + plaintext;

Self(a, b) } } +impl<const Q: u64, const N: usize, const K: usize> AddAssign for GLWE<Q, N, K> { + fn add_assign(&mut self, rhs: Self) { + for i in 0..K { + self.0 .0[i] = self.0 .0[i] + rhs.0 .0[i]; + } + self.1 = self.1 + rhs.1; + } +} +impl<const Q: u64, const N: usize, const K: usize> Sum<GLWE<Q, N, K>> for GLWE<Q, N, K> { + fn sum<I>(iter: I) -> Self + where + I: Iterator<Item = Self>, + { + let mut acc = GLWE::<Q, N, K>::zero(); + for e in iter { + acc += e; + } + acc + } +} + +impl<const Q: u64, const N: usize, const K: usize> Sub<GLWE<Q, N, K>> for GLWE<Q, N, K> { + type Output = Self; + fn sub(self, other: Self) -> Self { + let a: TR<Rq<Q, N>, K> = self.0 - other.0; + let b: Rq<Q, N> = self.1 - other.1; + Self(a, b) + } +} + impl<const Q: u64, const N: usize, const K: usize> Mul<Rq<Q, N>> for GLWE<Q, N, K> { type Output = Self; fn mul(self, plaintext: Rq<Q, N>) -> Self {

@@ -114,6 +187,15 @@ plaintext.compute_evals();

let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * plaintext).collect()); let b: Rq<Q, N> = self.1 * plaintext; + Self(a, b) + } +} + +impl<const Q: u64, const N: usize, const K: usize> Mul<Zq<Q>> for GLWE<Q, N, K> { + type Output = Self; + fn mul(self, e: Zq<Q>) -> Self { + let a: TR<Rq<Q, N>, K> = TR(self.0 .0.iter().map(|r_i| *r_i * e).collect()); + let b: Rq<Q, N> = self.1 * e; Self(a, b) } }

@@ -141,11 +223,18 @@ let (sk, pk) = S::new_key(&mut rng)?;

let msg_dist = Uniform::new(0_u64, T); let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; + let m: Rq<Q, N> = m.remodule::<Q>(); let c = S::encrypt(&mut rng, &pk, &m, delta)?; - let m_recovered = c.decrypt(&sk, delta); + let m_recovered = c.decrypt::<T>(&sk, delta); + + assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>()); + + // same but using encrypt_s (with sk instead of pk)) + let c = S::encrypt_s(&mut rng, &sk, &m, delta)?; + let m_recovered = c.decrypt::<T>(&sk, delta); - assert_eq!(m, m_recovered); + assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>()); } Ok(())

@@ -168,15 +257,17 @@

let msg_dist = Uniform::new(0_u64, T); let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; + let m1: Rq<Q, N> = m1.remodule::<Q>(); + let m2: Rq<Q, N> = m2.remodule::<Q>(); let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?; let c2 = S::encrypt(&mut rng, &pk, &m2, delta)?; let c3 = c1 + c2; - let m3_recovered = c3.decrypt(&sk, delta); + let m3_recovered = c3.decrypt::<T>(&sk, delta); - assert_eq!(m1 + m2, m3_recovered); + assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>()); } Ok(())

@@ -199,15 +290,17 @@

let msg_dist = Uniform::new(0_u64, T); let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; - let m2_scaled: Rq<Q, N> = m2.remodule::<Q>() * delta; + let m1: Rq<Q, N> = m1.remodule::<Q>(); + let m2: Rq<Q, N> = m2.remodule::<Q>(); + let m2_scaled: Rq<Q, N> = m2 * delta; let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?; let c3 = c1 + m2_scaled; - let m3_recovered = c3.decrypt(&sk, delta); + let m3_recovered = c3.decrypt::<T>(&sk, delta); - assert_eq!(m1 + m2, m3_recovered); + assert_eq!((m1 + m2).remodule::<T>(), m3_recovered.remodule::<T>()); } Ok(())

@@ -230,12 +323,14 @@

let msg_dist = Uniform::new(0_u64, T); let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; + let m1: Rq<Q, N> = m1.remodule::<Q>(); let m2: Rq<Q, N> = m2.remodule::<Q>(); let c1 = S::encrypt(&mut rng, &pk, &m1, delta)?; let c3 = c1 * m2; - let m3_recovered: Rq<T, N> = c3.decrypt(&sk, delta); + let m3_recovered: Rq<Q, N> = c3.decrypt::<T>(&sk, delta); + let m3_recovered: Rq<T, N> = m3_recovered.remodule::<T>(); assert_eq!((m1.to_r() * m2.to_r()).to_rq::<T>(), m3_recovered); }

@@ -265,18 +360,60 @@ let (sk, pk) = S::new_key(&mut rng)?;

let msg_dist = Uniform::new(0_u64, T); let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; + let m: Rq<Q, N> = m.remodule::<Q>(); let c = S::encrypt(&mut rng, &pk, &m, delta)?; + // let c = S::encrypt_s(&mut rng, &sk, &m, delta)?; let c2 = c.mod_switch::<P>(); let sk2: SecretKey<P, N, K> = SecretKey(TR(sk.0 .0.iter().map(|s_i| s_i.remodule::<P>()).collect())); let delta2: u64 = ((P as f64 * delta as f64) / Q as f64).round() as u64; - let m_recovered = c2.decrypt(&sk2, delta2); + let m_recovered = c2.decrypt::<T>(&sk2, delta2); - assert_eq!(m, m_recovered); + assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>()); } + + Ok(()) + } + + #[test] + fn test_key_switch() -> Result<()> { + const Q: u64 = 2u64.pow(16) + 1; + const N: usize = 128; + const T: u64 = 2; // plaintext modulus + const K: usize = 16; + type S = GLWE<Q, N, K>; + + let beta: u32 = 2; + let l: u32 = 16; + + let delta: u64 = Q / T; // floored + let mut rng = rand::thread_rng(); + + let (sk, pk) = S::new_key(&mut rng)?; + let (sk2, _) = S::new_key(&mut rng)?; + // ksk to switch from sk to sk2 + let ksk = S::new_ksk(&mut rng, beta, l, &sk, &sk2)?; + + let msg_dist = Uniform::new(0_u64, T); + let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?; + let m: Rq<Q, N> = m.remodule::<Q>(); + + let c = S::encrypt_s(&mut rng, &sk, &m, delta)?; + + let c2 = c.key_switch(beta, l, &ksk); + + // decrypt with the 2nd secret key + let m_recovered = c2.decrypt::<T>(&sk2, delta); + assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>()); + + // do the same but now encrypting with pk + // let c = S::encrypt(&mut rng, &pk, &m, delta)?; + // let c2 = c.key_switch(beta, l, &ksk); + // let m_recovered = c2.decrypt::<T>(&sk2, delta); + // assert_eq!(m.remodule::<T>(), m_recovered.remodule::<T>()); Ok(()) }
M gfhe/src/lib.rs → gfhe/src/lib.rs

@@ -5,4 +5,5 @@ #![allow(non_camel_case_types)]

#![allow(clippy::upper_case_acronyms)] #![allow(dead_code)] // TMP +pub mod glev; pub mod glwe;