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

add Ring trait, adapt R & Rq to it; add TR (tuple_ring)

Browse Source
This commit is contained in:
arnaucube
2025-07-09 16:44:39 +02:00
parent 84c54e8edd
commit 17b1e9ed43
8 changed files with 264 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
arith/src/lib.rs
View File

@@ -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;

74
arith/src/ring.rs
View File

@@ -1,14 +1,38 @@
//! 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 {
Self::from_vec_u64(rq.coeffs().to_vec().iter().map(|e| e.0).collect())
@@ -120,42 +144,72 @@ impl<const N: usize> PartialEq for R<N> {
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 @@ pub fn mod_centered_q<const Q: u64, const N: usize>(p: Vec<i128>) -> R<N> {
}
// 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 @@ impl<const N: usize> ops::Mul<&u64> for &R<N> {
}
}
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 {

97
arith/src/ringq.rs
View File

@@ -4,12 +4,15 @@
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.
#[derive(Clone, Copy)]
@@ -21,6 +24,28 @@ pub struct Rq<const Q: u64, const N: usize> {
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 @@ impl<const Q: u64, const N: usize> Rq<Q, 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 @@ impl<const Q: u64, const N: usize> PartialEq for Rq<Q, N> {
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 @@ impl<const Q: u64, const N: usize> ops::Add<Rq<Q, N>> for Rq<Q, N> {
// 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 @@ impl<const Q: u64, const N: usize> ops::Add<&Rq<Q, N>> for &Rq<Q, N> {
}
}
}
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 @@ impl<const Q: u64, const N: usize> ops::Sub<Rq<Q, N>> for Rq<Q, N> {
}
}
}
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 @@ impl<const Q: u64, const N: usize> ops::Sub<&Rq<Q, N>> for &Rq<Q, N> {
}
}
}
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 @@ impl<const Q: u64, const N: usize> ops::Mul<&Rq<Q, N>> for &Rq<Q, N> {
}
// 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 @@ impl<const Q: u64, const N: usize> ops::Mul<&Zq<Q>> for &Rq<Q, N> {
}
}
// 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 @@ impl<const Q: u64, const N: usize> ops::Mul<&u64> for &Rq<Q, N> {
}
}
// 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 @@ impl<const Q: u64, const N: usize> ops::Mul<&f64> for &Rq<Q, N> {
}
}
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 {

30
arith/src/traits.rs Normal file
View File

@@ -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
}

80
arith/src/tuple_ring.rs Normal file
View File

@@ -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())
}
}

8
arith/src/zq.rs
View File

@@ -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