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

37
arith/src/ntt.rs
View File

@@ -18,18 +18,12 @@ use std::sync::{Mutex, OnceLock};
static CACHE: OnceLock<Mutex<HashMap<(u64, usize), (Vec<Zq>, Vec<Zq>, Zq)>>> = OnceLock::new();
fn roots(q: u64, n: usize) -> (Vec<Zq>, Vec<Zq>, Zq) {
// Initialize CACHE with an empty HashMap on first use
let cache_lock = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
// Lock the HashMap for this thread
let mut cache = cache_lock.lock().unwrap();
if let Some(value) = cache.get(&(q, n)) {
// Found an existing value — return a clone
return value.clone();
}
// Not found — compute the new triple
let n_inv: Zq = Zq {
q,
v: const_inv_mod(q, n as u64),
@@ -37,10 +31,8 @@ fn roots(q: u64, n: usize) -> (Vec<Zq>, Vec<Zq>, Zq) {
let root_of_unity: u64 = primitive_root_of_unity(q, 2 * n);
let roots_of_unity: Vec<Zq> = roots_of_unity(q, n, root_of_unity);
let roots_of_unity_inv: Vec<Zq> = roots_of_unity_inv(q, n, roots_of_unity.clone());
let value = (roots_of_unity, roots_of_unity_inv, n_inv);
// Store and return
cache.insert((q, n), value.clone());
value
}
@@ -71,7 +63,8 @@ impl NTT {
t /= 2;
m *= 2;
}
// Rq::from_vec((a.q, n), r)
// TODO think if maybe not return a Rq type, or if returned Rq, maybe
// fill the `evals` field, which is what we're actually returning here
Rq {
param: RingParam { q, n },
coeffs: r,
@@ -107,10 +100,11 @@ impl NTT {
for i in 0..n {
r[i] = r[i] * n_inv;
}
// Rq::from_vec((a.q, n), r)
Rq {
param: RingParam { q, n },
coeffs: r,
// TODO maybe at `evals` place the inputed `a` which is the evals
// format
evals: None,
}
}
@@ -238,27 +232,4 @@ mod tests {
}
Ok(())
}
// #[test]
// fn test_ntt_loop_2() -> Result<()> {
// // let q: u64 = 2u64.pow(16) + 1;
// // let n: usize = 512;
// let q: u64 = 35184371138561;
// let n: usize = 1 << 14;
// let param = RingParam { q, n };
//
// use rand::distributions::Uniform;
// let mut rng = rand::thread_rng();
// let dist = Uniform::new(0_f64, q as f64);
//
// let a: Rq = Rq::rand(&mut rng, dist, &param);
// let start = std::time::Instant::now();
// for _ in 0..10_000 {
// let a_ntt = NTT::ntt(&a);
// let a_intt = NTT::intt(&a_ntt);
// assert_eq!(a, a_intt);
// }
// dbg!(start.elapsed());
// Ok(())
// }
}

4
arith/src/ring.rs
View File

@@ -33,10 +33,6 @@ pub trait Ring:
{
/// C defines the coefficient type
type C: Debug + Clone;
// type Param: Debug+Clone+Copy;
// const Q: u64;
// const N: usize;
fn param(&self) -> RingParam;
fn coeffs(&self) -> Vec<Self::C>;

22
arith/src/ring_n.rs
View File

@@ -15,14 +15,7 @@ pub struct R {
pub coeffs: Vec<i64>,
}
// impl<const N: usize> Ring for R<N> {
impl R {
// type C = i64;
// type Param = usize; // n
// const Q: u64 = i64::MAX as u64; // WIP
// const N: usize = N;
pub fn coeffs(&self) -> Vec<i64> {
self.coeffs.clone()
}
@@ -33,16 +26,12 @@ impl R {
}
}
fn rand(mut rng: impl Rng, dist: impl Distribution<f64>, n: usize) -> Self {
// let coeffs: [i64; N] = array::from_fn(|_| Self::C::rand(&mut rng, &dist));
// let coeffs: [i64; N] = array::from_fn(|_| dist.sample(&mut rng).round() as i64);
Self {
n,
coeffs: std::iter::repeat_with(|| dist.sample(&mut rng).round() as i64)
.take(n)
.collect(),
}
// let coeffs: [C; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng)));
// Self(coeffs)
}
pub fn from_vec(n: usize, coeffs: Vec<i64>) -> Self {
@@ -90,9 +79,6 @@ impl From<crate::ring_nq::Rq> for R {
}
impl R {
// pub fn coeffs(&self) -> [i64; N] {
// self.0
// }
pub fn to_rq(self, q: u64) -> crate::Rq {
crate::Rq::from((q, self))
}
@@ -196,7 +182,6 @@ impl Add<&R> for &R {
type Output = R;
fn add(self, rhs: &R) -> Self::Output {
// R(array::from_fn(|i| self.0[i] + rhs.0[i]))
assert_eq!(self.n, rhs.n);
R {
n: self.n,
@@ -220,11 +205,6 @@ impl Sum<R> for R {
where
I: Iterator<Item = Self>,
{
// let mut acc = R::zero();
// for e in iter {
// acc += e;
// }
// acc
let first = iter.next().unwrap();
iter.fold(first, |acc, x| acc + x)
}
@@ -234,7 +214,6 @@ impl Sub<R> for R {
type Output = Self;
fn sub(self, rhs: Self) -> Self {
// Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
assert_eq!(self.n, rhs.n);
Self {
n: self.n,
@@ -248,7 +227,6 @@ impl Sub<&R> for &R {
type Output = R;
fn sub(self, rhs: &R) -> Self::Output {
// R(array::from_fn(|i| self.0[i] - rhs.0[i]))
assert_eq!(self.n, rhs.n);
R {
n: self.n,

73
arith/src/ring_nq.rs
View File

@@ -13,9 +13,6 @@ use crate::zq::{modulus_u64, Zq};
use crate::{Ring, RingParam};
// NOTE: currently using fixed-size arrays, but pending to see if with
// real-world parameters the stack can keep up; if not will move everything to
// use Vec.
/// PolynomialRing element, where the PolynomialRing is R = Z_q[X]/(X^n +1)
/// The implementation assumes that q is prime.
#[derive(Clone)]
@@ -31,8 +28,6 @@ pub struct Rq {
impl Ring for Rq {
type C = Zq;
// type Param = (u64, usize);
// type Param = Param;
fn param(&self) -> RingParam {
self.param
@@ -40,8 +35,6 @@ impl Ring for Rq {
fn coeffs(&self) -> Vec<Self::C> {
self.coeffs.to_vec()
}
// fn zero(q: u64, n: usize) -> Self {
// fn zero(param: (u64, usize)) -> Self {
fn zero(param: &RingParam) -> Self {
Self {
param: param.clone(),
@@ -50,8 +43,6 @@ impl Ring for Rq {
}
}
fn rand(mut rng: impl Rng, dist: impl Distribution<f64>, param: &RingParam) -> Self {
// let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng)));
// let coeffs: [Zq<Q>; N] = array::from_fn(|_| Self::C::rand(&mut rng, &dist));
Self {
param: param.clone(),
coeffs: std::iter::repeat_with(|| Self::C::rand(&mut rng, &dist, param.q))
@@ -93,23 +84,17 @@ impl Ring for Rq {
q: p,
n: self.param.n,
};
// Rq::from_vec_u64(p, self.n, self.coeffs().iter().map(|m_i| m_i.v).collect())
Rq::from_vec_u64(&param, self.coeffs().iter().map(|m_i| m_i.v).collect())
}
/// perform the mod switch operation from Q to Q', where Q2=Q'
// fn mod_switch<const P: u64, const M: usize>(&self) -> impl Ring {
fn mod_switch(&self, p: u64) -> Rq {
let param = RingParam {
q: p,
n: self.param.n,
};
// assert_eq!(N, M); // sanity check
Rq {
param,
// q: p,
// n: self.n,
// coeffs: array::from_fn(|i| self.coeffs[i].mod_switch::<P>()),
coeffs: self.coeffs.iter().map(|c_i| c_i.mod_switch(p)).collect(),
evals: None,
}
@@ -135,7 +120,6 @@ impl From<(u64, crate::ring_n::R)> for Rq {
Self::from_vec(
&RingParam { q, n: r.n },
// (q, r.n),
r.coeffs()
.iter()
.map(|e| Zq::from_f64(q, *e as f64))
@@ -161,21 +145,13 @@ impl Rq {
self.coeffs.clone()
}
pub fn compute_evals(&mut self) {
self.evals = Some(NTT::ntt(self).coeffs); // TODO improve, ntt returns Rq but here
// just needs Vec<Zq>
self.evals = Some(NTT::ntt(self).coeffs);
// TODO improve, ntt returns Rq but here just needs Vec<Zq>
}
pub fn to_r(self) -> crate::R {
crate::R::from(self)
}
// TODO rm since it is implemented in Ring trait impl
// pub fn zero() -> Self {
// let coeffs = array::from_fn(|_| Zq::zero());
// Self {
// coeffs,
// evals: None,
// }
// }
// this method is mostly for tests
pub fn from_vec_u64(param: &RingParam, coeffs: Vec<u64>) -> Self {
let coeffs_mod_q: Vec<Zq> = coeffs.iter().map(|c| Zq::from_u64(param.q, *c)).collect();
@@ -204,14 +180,9 @@ impl Rq {
mut rng: impl Rng,
dist: impl Distribution<f64>,
param: &RingParam,
// q: u64,
// n: usize,
) -> Result<Self> {
// let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
Ok(Self {
param: *param,
// q,
// n,
coeffs: std::iter::repeat_with(|| Zq::from_f64(param.q, dist.sample(&mut rng).abs()))
.take(param.n)
.collect(),
@@ -222,14 +193,9 @@ impl Rq {
mut rng: impl Rng,
dist: impl Distribution<f64>,
param: &RingParam,
// q: u64,
// n: usize,
) -> Result<Self> {
// let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
Ok(Self {
param: *param,
// q,
// n,
coeffs: std::iter::repeat_with(|| Zq::from_f64(param.q, dist.sample(&mut rng).abs()))
.take(param.n)
.collect(),
@@ -241,7 +207,6 @@ impl Rq {
dist: impl Distribution<f64>,
param: &RingParam,
) -> Result<Self> {
// let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng)));
Ok(Self {
param: *param,
coeffs: std::iter::repeat_with(|| Zq::from_f64(param.q, dist.sample(&mut rng)))
@@ -255,7 +220,6 @@ impl Rq {
dist: impl Distribution<u64>,
param: &RingParam,
) -> Result<Self> {
// let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng)));
Ok(Self {
param: *param,
coeffs: std::iter::repeat_with(|| Zq::from_u64(param.q, dist.sample(&mut rng)))
@@ -270,7 +234,6 @@ impl Rq {
dist: impl Distribution<bool>,
param: &RingParam,
) -> Result<Self> {
// let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_bool(dist.sample(&mut rng)));
Ok(Rq {
param: *param,
coeffs: std::iter::repeat_with(|| Zq::from_bool(param.q, dist.sample(&mut rng)))
@@ -304,7 +267,6 @@ impl Rq {
pub fn mul_by_zq(&self, s: &Zq) -> Self {
Self {
param: self.param,
// coeffs: array::from_fn(|i| self.coeffs[i] * *s),
coeffs: self.coeffs.iter().map(|c_i| *c_i * *s).collect(),
evals: None,
}
@@ -313,7 +275,6 @@ impl Rq {
let s = Zq::from_u64(self.param.q, s);
Self {
param: self.param,
// coeffs: array::from_fn(|i| self.coeffs[i] * s),
coeffs: self.coeffs.iter().map(|&e| e * s).collect(),
evals: None,
}
@@ -321,7 +282,6 @@ impl Rq {
pub fn mul_by_f64(&self, s: f64) -> Self {
Self {
param: self.param,
// coeffs: array::from_fn(|i| Zq::from_f64(self.coeffs[i].0 as f64 * s)),
coeffs: self
.coeffs
.iter()
@@ -450,22 +410,11 @@ impl Add<Rq> for Rq {
assert_eq!(self.param, rhs.param);
Self {
param: self.param,
// coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
coeffs: zip_eq(self.coeffs, rhs.coeffs)
.map(|(l, r)| l + r)
.collect(),
evals: None,
}
// Self {
// coeffs: self
// .coeffs
// .iter()
// .zip(rhs.coeffs)
// .map(|(a, b)| *a + b)
// .collect(),
// evals: None,
// }
// Self(r.iter_mut().map(|e| e.r#mod()).collect()) // TODO mod should happen auto in +
}
}
impl Add<&Rq> for &Rq {
@@ -475,7 +424,6 @@ impl Add<&Rq> for &Rq {
assert_eq!(self.param, rhs.param);
Rq {
param: self.param,
// coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
coeffs: zip_eq(self.coeffs.clone(), rhs.coeffs.clone())
.map(|(l, r)| l + r)
.collect(),
@@ -497,11 +445,6 @@ impl Sum<Rq> for Rq {
where
I: Iterator<Item = Self>,
{
// let mut acc = Rq::zero();
// for e in iter {
// acc += e;
// }
// acc
let first = iter.next().unwrap();
iter.fold(first, |acc, x| acc + x)
}
@@ -514,7 +457,6 @@ impl Sub<Rq> for Rq {
assert_eq!(self.param, rhs.param);
Self {
param: self.param,
// coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
coeffs: zip_eq(self.coeffs, rhs.coeffs)
.map(|(l, r)| l - r)
.collect(),
@@ -529,7 +471,6 @@ impl Sub<&Rq> for &Rq {
debug_assert_eq!(self.param, rhs.param);
Rq {
param: self.param,
// coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
coeffs: zip_eq(self.coeffs.clone(), rhs.coeffs.clone())
.map(|(l, r)| l - r)
.collect(),
@@ -613,8 +554,6 @@ impl Neg for Rq {
fn neg(self) -> Self::Output {
Self {
param: self.param,
// coeffs: array::from_fn(|i| -self.coeffs[i]),
// coeffs: self.coeffs.iter().map(|c_i| -c_i).collect(),
coeffs: self.coeffs.iter().map(|c_i| -*c_i).collect(),
evals: None,
}
@@ -624,7 +563,6 @@ impl Neg for Rq {
// note: this assumes that Q is prime
fn mul_mut(lhs: &mut Rq, rhs: &mut Rq) -> Rq {
assert_eq!(lhs.param, rhs.param);
// let (q, n) = (lhs.q, lhs.n);
// reuse evaluations if already computed
if !lhs.evals.is_some() {
@@ -636,7 +574,6 @@ fn mul_mut(lhs: &mut Rq, rhs: &mut Rq) -> Rq {
let lhs_evals = lhs.evals.clone().unwrap();
let rhs_evals = rhs.evals.clone().unwrap();
// let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
let c_ntt: Rq = Rq::from_vec(
&lhs.param,
zip_eq(lhs_evals, rhs_evals).map(|(l, r)| l * r).collect(),
@@ -645,7 +582,7 @@ fn mul_mut(lhs: &mut Rq, rhs: &mut Rq) -> Rq {
Rq::new(&lhs.param, c.coeffs, Some(c_ntt.coeffs))
}
// note: this assumes that Q is prime
// TODO impl karatsuba for non-prime Q
// TODO impl karatsuba for non-prime Q. Alternatively check NTT with RNS trick.
fn mul(lhs: &Rq, rhs: &Rq) -> Rq {
assert_eq!(lhs.param, rhs.param);
@@ -661,7 +598,6 @@ fn mul(lhs: &Rq, rhs: &Rq) -> Rq {
NTT::ntt(rhs).coeffs
};
// let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
let c_ntt: Rq = Rq::from_vec(
&lhs.param,
zip_eq(lhs_evals, rhs_evals).map(|(l, r)| l * r).collect(),
@@ -758,11 +694,8 @@ mod tests {
assert_eq!(a.len(), param.n);
assert_eq!(b.len(), param.n);
// let a: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(a[i]));
let mut a = Rq::from_vec_u64(&param, a);
// let b: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(b[i]));
let mut b = Rq::from_vec_u64(&param, b);
// let expected_c: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(expected_c[i]));
let expected_c = Rq::from_vec_u64(&param, expected_c);
let c = mul_mut(&mut a, &mut b);

22
arith/src/ring_torus.rs
View File

@@ -22,17 +22,12 @@ use crate::{
/// 𝕋, where Q=2^64.
#[derive(Clone, Debug)]
pub struct Tn {
// pub n: usize,
pub param: RingParam,
pub coeffs: Vec<T64>,
}
impl Ring for Tn {
type C = T64;
// type Param = usize; // n
// const Q: u64 = u64::MAX; // WIP
// const N: usize = N;
fn param(&self) -> RingParam {
RingParam {
@@ -58,7 +53,6 @@ impl Ring for Tn {
.take(param.n)
.collect(),
}
// Self(array::from_fn(|_| T64::rand(&mut rng, &dist)))
}
fn from_vec(param: &RingParam, coeffs: Vec<Self::C>) -> Self {
@@ -160,7 +154,6 @@ impl Add<Tn> for Tn {
type Output = Self;
fn add(self, rhs: Self) -> Self {
// Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
assert_eq!(self.param, rhs.param);
Self {
param: self.param,
@@ -174,7 +167,6 @@ impl Add<&Tn> for &Tn {
type Output = Tn;
fn add(self, rhs: &Tn) -> Self::Output {
// Tn(array::from_fn(|i| self.0[i] + rhs.0[i]))
assert_eq!(self.param, rhs.param);
Tn {
param: self.param,
@@ -198,11 +190,6 @@ impl Sum<Tn> for Tn {
where
I: Iterator<Item = Self>,
{
// let mut acc = Tn::<N>::zero();
// for e in iter {
// acc += e;
// }
// acc
let first = iter.next().unwrap();
iter.fold(first, |acc, x| acc + x)
}
@@ -225,7 +212,6 @@ impl Sub<&Tn> for &Tn {
type Output = Tn;
fn sub(self, rhs: &Tn) -> Self::Output {
// Tn(array::from_fn(|i| self.0[i] - rhs.0[i]))
assert_eq!(self.param, rhs.param);
Tn {
param: self.param,
@@ -238,9 +224,6 @@ impl Sub<&Tn> for &Tn {
impl SubAssign for Tn {
fn sub_assign(&mut self, rhs: Self) {
// for i in 0..N {
// self.0[i] -= rhs.0[i];
// }
assert_eq!(self.param, rhs.param);
for i in 0..self.param.n {
self.coeffs[i] -= rhs.coeffs[i];
@@ -252,7 +235,6 @@ impl Neg for Tn {
type Output = Self;
fn neg(self) -> Self::Output {
// Tn(array::from_fn(|i| -self.0[i]))
Self {
param: self.param,
coeffs: self.coeffs.iter().map(|c_i| -*c_i).collect(),
@@ -300,7 +282,6 @@ fn naive_poly_mul(poly1: &Tn, poly2: &Tn) -> Tn {
Tn {
param,
// coeffs: array::from_fn(|i| T64(result[i] as u64)),
coeffs: result.iter().map(|r_i| T64(*r_i as u64)).collect(),
}
}
@@ -321,7 +302,6 @@ impl Mul<T64> for Tn {
fn mul(self, s: T64) -> Self {
Self {
param: self.param,
// coeffs: array::from_fn(|i| self.coeffs[i] * s),
coeffs: self.coeffs.iter().map(|c_i| *c_i * s).collect(),
}
}
@@ -330,7 +310,6 @@ impl Mul<T64> for Tn {
impl Mul<u64> for Tn {
type Output = Self;
fn mul(self, s: u64) -> Self {
// Self(array::from_fn(|i| self.0[i] * s))
Tn {
param: self.param,
coeffs: self.coeffs.iter().map(|c_i| *c_i * s).collect(),
@@ -340,7 +319,6 @@ impl Mul<u64> for Tn {
impl Mul<&u64> for &Tn {
type Output = Tn;
fn mul(self, s: &u64) -> Self::Output {
// Tn::<N>(array::from_fn(|i| self.0[i] * *s))
Tn {
param: self.param,
coeffs: self.coeffs.iter().map(|c_i| c_i * s).collect(),

12
arith/src/torus.rs
View File

@@ -16,10 +16,6 @@ pub struct T64(pub u64);
// `Tn<1>`.
impl Ring for T64 {
type C = T64;
// type Param = ();
// const Q: u64 = u64::MAX; // WIP
// const N: usize = 1;
fn param(&self) -> RingParam {
RingParam {
@@ -45,13 +41,13 @@ impl Ring for T64 {
// TODO rm beta & l from inputs, make it always beta=2,l=64.
/// Note: only beta=2 and l=64 is supported.
fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
assert_eq!(beta, 2u32); // only beta=2 supported
// assert_eq!(l, 64u32); // only l=64 supported
assert_eq!(beta, 2u32, "only beta=2 supported");
// assert_eq!(l, 64u32, "only l=64 supported");
// (0..64)
(0..l)
(0..l as u64)
.rev()
.map(|i| T64(((self.0 >> i) & 1) as u64))
.map(|i| T64((self.0 >> i) & 1))
.collect()
}
fn remodule(&self, p: u64) -> T64 {

10
arith/src/zq.rs
View File

@@ -9,13 +9,6 @@ pub struct Zq {
pub v: u64,
}
// WIP
// impl<const Q: u64> From<Vec<u64>> for Vec<Zq<Q>> {
// fn from(v: Vec<u64>) -> Self {
// v.into_iter().map(Zq::new).collect()
// }
// }
pub(crate) fn modulus_u64(q: u64, e: u64) -> u64 {
(e % q + q) % q
}
@@ -24,7 +17,6 @@ impl Zq {
// TODO WIP
let r: f64 = dist.sample(&mut rng);
Self::from_f64(q, r)
// Self::from_u64(r.round() as u64)
}
pub fn from_u64(q: u64, v: u64) -> Self {
if v >= q {
@@ -81,7 +73,7 @@ impl Zq {
// for rem != Self(0) {
while rem != Self::zero(self.q) {
// if odd
// TODO use a more readible expression
// TODO use a more readeable expression
if 1 - ((rem.v & 1) << 1) as i64 == -1 {
res = res * exp;
}