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

arith/Cargo.toml

@@ -7,6 +7,7 @@ edition = "2024"
 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

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;

arith/src/ring.rs

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

arith/src/ringq.rs

@@ -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 {
+    // TODO rm since it is implemented in Ring trait impl
-        let coeffs = array::from_fn(|_| Zq::zero());
+    // pub fn zero() -> Self {
-        Self {
+    //     let coeffs = array::from_fn(|_| Zq::zero());
-            coeffs,
+    //     Self {
-            evals: None,
+    //         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 {

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

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

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

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