arnaucube
arnaucube/fhe-studyPublic Self-hosted repository
add Ring trait, adapt R & Rq to it; add TR (tuple_ring)
arnaucube git@arnaucube.com
Wed, 09 Jul 2025 16:44:39 +0200
commit

17b1e9ed43d2a1c93bf4027de9d54411ee26fed2

parent

84c54e8edd47b1755063a016fca70e1baf6a3a23

M arith/Cargo.toml → arith/Cargo.toml

@@ -7,6 +7,7 @@ [dependencies]

anyhow = { workspace = true } rand = { workspace = true } rand_distr = { workspace = true } +itertools = { workspace = true } # TMP: the next 4 imports are TMP, to solve systems of linear equations. Used # for the CKKS encoding step, probably remvoed once in ckks the encoding is done
M arith/src/lib.rs → arith/src/lib.rs

@@ -10,6 +10,8 @@ mod naive_ntt; // note: for dev only

pub mod ntt; pub mod ring; pub mod ringq; +pub mod traits; +pub mod tuple_ring; pub mod zq; pub use complex::C;

@@ -17,4 +19,6 @@ pub use matrix::Matrix;

pub use ntt::NTT; pub use ring::R; pub use ringq::Rq; +pub use traits::Ring; +pub use tuple_ring::TR; pub use zq::Zq;
M arith/src/ring.rs → arith/src/ring.rs

@@ -1,13 +1,37 @@

//! Polynomial ring Z[X]/(X^N+1) //! +use anyhow::Result; +use rand::{distributions::Distribution, Rng}; use std::array; use std::fmt; -use std::ops; +use std::iter::Sum; +use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign}; + +use crate::Ring; +// TODO rename to not have name conflicts with the Ring trait (R: Ring) // PolynomialRing element, where the PolynomialRing is R = Z[X]/(X^n +1) #[derive(Clone, Copy)] pub struct R<const N: usize>(pub [i64; N]); + +impl<const N: usize> Ring for R<N> { + type C = i64; + fn coeffs(&self) -> Vec<Self::C> { + self.0.to_vec() + } + fn zero() -> Self { + let coeffs: [i64; N] = array::from_fn(|_| 0i64); + Self(coeffs) + } + fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> 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(coeffs) + // let coeffs: [C; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng))); + // Self(coeffs) + } +} impl<const Q: u64, const N: usize> From<crate::ringq::Rq<Q, N>> for R<N> { fn from(rq: crate::ringq::Rq<Q, N>) -> Self {

@@ -120,42 +144,72 @@ fn eq(&self, other: &Self) -> bool {

self.0 == other.0 } } -impl<const N: usize> ops::Add<R<N>> for R<N> { +impl<const N: usize> Add<R<N>> for R<N> { type Output = Self; fn add(self, rhs: Self) -> Self { Self(array::from_fn(|i| self.0[i] + rhs.0[i])) } } -impl<const N: usize> ops::Add<&R<N>> for &R<N> { +impl<const N: usize> Add<&R<N>> for &R<N> { type Output = R<N>; fn add(self, rhs: &R<N>) -> Self::Output { R(array::from_fn(|i| self.0[i] + rhs.0[i])) } } -impl<const N: usize> ops::Sub<R<N>> for R<N> { +impl<const N: usize> AddAssign for R<N> { + fn add_assign(&mut self, rhs: Self) { + for i in 0..N { + self.0[i] += rhs.0[i]; + } + } +} + +impl<const N: usize> Sum<R<N>> for R<N> { + fn sum<I>(iter: I) -> Self + where + I: Iterator<Item = Self>, + { + let mut acc = R::<N>::zero(); + for e in iter { + acc += e; + } + acc + } +} + +impl<const N: usize> Sub<R<N>> for R<N> { type Output = Self; fn sub(self, rhs: Self) -> Self { Self(array::from_fn(|i| self.0[i] - rhs.0[i])) } } -impl<const N: usize> ops::Sub<&R<N>> for &R<N> { +impl<const N: usize> Sub<&R<N>> for &R<N> { type Output = R<N>; fn sub(self, rhs: &R<N>) -> Self::Output { R(array::from_fn(|i| self.0[i] - rhs.0[i])) } } -impl<const N: usize> ops::Mul<R<N>> for R<N> { + +impl<const N: usize> SubAssign for R<N> { + fn sub_assign(&mut self, rhs: Self) { + for i in 0..N { + self.0[i] -= rhs.0[i]; + } + } +} + +impl<const N: usize> Mul<R<N>> for R<N> { type Output = Self; fn mul(self, rhs: Self) -> Self { naive_poly_mul(&self, &rhs) } } -impl<const N: usize> ops::Mul<&R<N>> for &R<N> { +impl<const N: usize> Mul<&R<N>> for &R<N> { type Output = R<N>; fn mul(self, rhs: &R<N>) -> Self::Output {

@@ -255,14 +309,14 @@ R::<N>::from_vec(r.iter().map(|v| *v as i64).collect::<Vec<i64>>())

} // mul by u64 -impl<const N: usize> ops::Mul<u64> for R<N> { +impl<const N: usize> Mul<u64> for R<N> { type Output = Self; fn mul(self, s: u64) -> Self { self.mul_by_i64(s as i64) } } -impl<const N: usize> ops::Mul<&u64> for &R<N> { +impl<const N: usize> Mul<&u64> for &R<N> { type Output = R<N>; fn mul(self, s: &u64) -> Self::Output {

@@ -270,7 +324,7 @@ self.mul_by_i64(*s as i64)

} } -impl<const N: usize> ops::Neg for R<N> { +impl<const N: usize> Neg for R<N> { type Output = Self; fn neg(self) -> Self::Output {
M arith/src/ringq.rs → arith/src/ringq.rs

@@ -4,11 +4,14 @@

use rand::{distributions::Distribution, Rng}; use std::array; use std::fmt; -use std::ops; +use std::iter::Sum; +use std::ops::{Add, AddAssign, Mul, Neg, Sub, SubAssign}; use crate::ntt::NTT; use crate::zq::{modulus_u64, Zq}; use anyhow::{anyhow, Result}; + +use crate::Ring; /// PolynomialRing element, where the PolynomialRing is R = Z_q[X]/(X^n +1) /// The implementation assumes that q is prime.

@@ -21,6 +24,28 @@ // multiplications avoiding recomputing it

pub(crate) evals: Option<[Zq<Q>; N]>, } +impl<const Q: u64, const N: usize> Ring for Rq<Q, N> { + type C = Zq<Q>; + fn coeffs(&self) -> Vec<Self::C> { + self.coeffs.to_vec() + } + fn zero() -> Self { + let coeffs = array::from_fn(|_| Zq::zero()); + Self { + coeffs, + evals: None, + } + } + fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> 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 { + coeffs, + evals: None, + } + } +} + // TODO define a trait "PolynomialRingTrait" or similar, so that when other structs use it can just // use the trait and not need to add '<Q, N>' to their params

@@ -59,13 +84,14 @@ pub fn to_r(self) -> crate::R<N> {

crate::R::<N>::from(self) } - pub fn zero() -> Self { - let coeffs = array::from_fn(|_| Zq::zero()); - Self { - coeffs, - evals: None, - } - } + // 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, + // } + // } pub fn from_vec(coeffs: Vec<Zq<Q>>) -> Self { let mut p = coeffs; modulus::<Q, N>(&mut p);

@@ -292,7 +318,7 @@ fn eq(&self, other: &Self) -> bool {

self.coeffs == other.coeffs } } -impl<const Q: u64, const N: usize> ops::Add<Rq<Q, N>> for Rq<Q, N> { +impl<const Q: u64, const N: usize> Add<Rq<Q, N>> for Rq<Q, N> { type Output = Self; fn add(self, rhs: Self) -> Self {

@@ -312,7 +338,7 @@ // }

// Self(r.iter_mut().map(|e| e.r#mod()).collect()) // TODO mod should happen auto in + } } -impl<const Q: u64, const N: usize> ops::Add<&Rq<Q, N>> for &Rq<Q, N> { +impl<const Q: u64, const N: usize> Add<&Rq<Q, N>> for &Rq<Q, N> { type Output = Rq<Q, N>; fn add(self, rhs: &Rq<Q, N>) -> Self::Output {

@@ -322,7 +348,28 @@ evals: None,

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

@@ -332,7 +379,7 @@ evals: None,

} } } -impl<const Q: u64, const N: usize> ops::Sub<&Rq<Q, N>> for &Rq<Q, N> { +impl<const Q: u64, const N: usize> Sub<&Rq<Q, N>> for &Rq<Q, N> { type Output = Rq<Q, N>; fn sub(self, rhs: &Rq<Q, N>) -> Self::Output {

@@ -342,14 +389,22 @@ evals: None,

} } } -impl<const Q: u64, const N: usize> ops::Mul<Rq<Q, N>> for Rq<Q, N> { +impl<const Q: u64, const N: usize> SubAssign for Rq<Q, N> { + fn sub_assign(&mut self, rhs: Self) { + for i in 0..N { + self.coeffs[i] -= rhs.coeffs[i]; + } + } +} + +impl<const Q: u64, const N: usize> Mul<Rq<Q, N>> for Rq<Q, N> { type Output = Self; fn mul(self, rhs: Self) -> Self { mul(&self, &rhs) } } -impl<const Q: u64, const N: usize> ops::Mul<&Rq<Q, N>> for &Rq<Q, N> { +impl<const Q: u64, const N: usize> Mul<&Rq<Q, N>> for &Rq<Q, N> { type Output = Rq<Q, N>; fn mul(self, rhs: &Rq<Q, N>) -> Self::Output {

@@ -358,14 +413,14 @@ }

} // mul by Zq element -impl<const Q: u64, const N: usize> ops::Mul<Zq<Q>> for Rq<Q, N> { +impl<const Q: u64, const N: usize> Mul<Zq<Q>> for Rq<Q, N> { type Output = Self; fn mul(self, s: Zq<Q>) -> Self { self.mul_by_zq(&s) } } -impl<const Q: u64, const N: usize> ops::Mul<&Zq<Q>> for &Rq<Q, N> { +impl<const Q: u64, const N: usize> Mul<&Zq<Q>> for &Rq<Q, N> { type Output = Rq<Q, N>; fn mul(self, s: &Zq<Q>) -> Self::Output {

@@ -373,14 +428,14 @@ self.mul_by_zq(s)

} } // mul by u64 -impl<const Q: u64, const N: usize> ops::Mul<u64> for Rq<Q, N> { +impl<const Q: u64, const N: usize> Mul<u64> for Rq<Q, N> { type Output = Self; fn mul(self, s: u64) -> Self { self.mul_by_u64(s) } } -impl<const Q: u64, const N: usize> ops::Mul<&u64> for &Rq<Q, N> { +impl<const Q: u64, const N: usize> Mul<&u64> for &Rq<Q, N> { type Output = Rq<Q, N>; fn mul(self, s: &u64) -> Self::Output {

@@ -388,14 +443,14 @@ self.mul_by_u64(*s)

} } // mul by f64 -impl<const Q: u64, const N: usize> ops::Mul<f64> for Rq<Q, N> { +impl<const Q: u64, const N: usize> Mul<f64> for Rq<Q, N> { type Output = Self; fn mul(self, s: f64) -> Self { self.mul_by_f64(s) } } -impl<const Q: u64, const N: usize> ops::Mul<&f64> for &Rq<Q, N> { +impl<const Q: u64, const N: usize> Mul<&f64> for &Rq<Q, N> { type Output = Rq<Q, N>; fn mul(self, s: &f64) -> Self::Output {

@@ -403,7 +458,7 @@ self.mul_by_f64(*s)

} } -impl<const Q: u64, const N: usize> ops::Neg for Rq<Q, N> { +impl<const Q: u64, const N: usize> Neg for Rq<Q, N> { type Output = Self; fn neg(self) -> Self::Output {
A arith/src/traits.rs

@@ -0,0 +1,30 @@

+use anyhow::Result; +use rand::{distributions::Distribution, Rng}; +use std::fmt::Debug; +use std::iter::Sum; +use std::ops::{Add, AddAssign, Mul, Sub, SubAssign}; + +/// represents a ring element +pub trait Ring: + Sized + + Add<Output = Self> + + AddAssign + + Sum + + Sub<Output = Self> + + SubAssign + + Mul<Output = Self> + + Mul<u64, Output = Self> // scalar mul + + PartialEq + + Debug + + Clone + + Sum<<Self as Add>::Output> + + Sum<<Self as Mul>::Output> +{ + /// C defines the coefficient type + type C: Debug + Clone; + + fn coeffs(&self) -> Vec<Self::C>; + fn zero() -> Self; + fn rand(rng: impl Rng, dist: impl Distribution<f64>) -> Self; + // note/wip/warning: dist (0,q) with f64, will output more '0=q' elements than other values +}
A arith/src/tuple_ring.rs

@@ -0,0 +1,80 @@

+//! This file implements the struct for an Tuple of Ring Rq elements and its +//! operations. + +use anyhow::Result; +use itertools::zip_eq; +use rand::{distributions::Distribution, Rng}; +use rand_distr::{Normal, Uniform}; +use std::iter::Sum; +use std::{array, ops}; + +use crate::Ring; + +// #[derive(Clone, Copy, Debug)] +// pub struct TR<R: Ring, const K: usize>([R; K]); + +/// Tuple of K Ring (Rq) elements. We use Vec<R> to allocate it in the heap, +/// since if using a fixed-size array it would overflow the stack. +#[derive(Clone, Debug)] +pub struct TR<R: Ring, const K: usize>(Vec<R>); + +impl<R: Ring, const K: usize> TR<R, K> { + pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self { + Self( + (0..K) + .into_iter() + .map(|_| R::rand(&mut rng, &dist)) + .collect(), + ) + } +} + +impl<R: Ring, const K: usize> ops::Add<TR<R, K>> for TR<R, K> { + type Output = Self; + fn add(self, other: Self) -> Self { + Self( + zip_eq(self.0, other.0) + .map(|(s, o)| s + o) + .collect::<Vec<_>>(), + ) + } +} + +impl<R: Ring, const K: usize> ops::Sub<TR<R, K>> for TR<R, K> { + type Output = Self; + fn sub(self, other: Self) -> Self { + Self(zip_eq(self.0, other.0).map(|(s, o)| s - o).collect()) + } +} + +/// for (TR,TR), the Mul operation is defined as: +/// for A, B \in R^k, result = Σ A_i * B_i \in R +impl<R: Ring, const K: usize> ops::Mul<TR<R, K>> for TR<R, K> { + type Output = R; + fn mul(self, other: Self) -> R { + zip_eq(self.0, other.0).map(|(s, o)| s * o).sum() + } +} +impl<R: Ring, const K: usize> ops::Mul<&TR<R, K>> for &TR<R, K> { + type Output = R; + fn mul(self, other: &TR<R, K>) -> R { + zip_eq(self.0.clone(), other.0.clone()) + .map(|(s, o)| s * o) + .sum() + } +} + +/// for (TR, R), the Mul operation is defined as each element of TR is +/// multiplied by R +impl<R: Ring, const K: usize> ops::Mul<R> for TR<R, K> { + type Output = TR<R, K>; + fn mul(self, other: R) -> TR<R, K> { + Self(self.0.iter().map(|s| s.clone() * other.clone()).collect()) + } +} +impl<R: Ring, const K: usize> ops::Mul<&R> for &TR<R, K> { + type Output = TR<R, K>; + fn mul(self, other: &R) -> TR<R, K> { + TR::<R, K>(self.0.iter().map(|s| s.clone() * other.clone()).collect()) + } +}
M arith/src/zq.rs → arith/src/zq.rs

@@ -1,3 +1,5 @@

+use anyhow::{anyhow, Result}; +use rand::{distributions::Distribution, Rng}; use std::fmt; use std::ops;

@@ -16,6 +18,12 @@ pub(crate) fn modulus_u64<const Q: u64>(e: u64) -> u64 {

(e % Q + Q) % Q } impl<const Q: u64> Zq<Q> { + pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self { + // TODO WIP + let r: f64 = dist.sample(&mut rng); + Self::from_f64(r) + // Self::from_u64(r.round() as u64) + } pub fn from_u64(e: u64) -> Self { if e >= Q { // (e % Q + Q) % Q
M bfv/src/lib.rs → bfv/src/lib.rs

@@ -10,7 +10,7 @@ use rand::Rng;

use rand_distr::{Normal, Uniform}; use std::ops; -use arith::{Rq, R}; +use arith::{Ring, Rq, R}; // error deviation for the Gaussian(Normal) distribution // sigma=3.2 from: https://eprint.iacr.org/2022/162.pdf page 5