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 wip version of tensor & relinearization

Browse Source
This commit is contained in:
arnaucube
2025-06-22 15:51:20 +02:00
parent f3a368ab6a
commit d2fc32ac0c
10 changed files with 1145 additions and 469 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
arithmetic/src/lib.rs
View File

@@ -4,11 +4,13 @@
#![allow(clippy::upper_case_acronyms)]
#![allow(dead_code)] // TMP
mod naive; // TODO rm
mod naive_ntt; // TODO rm
pub mod ntt;
pub mod ring;
pub mod ringq;
pub mod zq;
pub use ntt::NTT;
pub use ring::PR;
pub use ring::R;
pub use ringq::Rq;
pub use zq::Zq;

10
arithmetic/src/naive.rs → arithmetic/src/naive_ntt.rs
View File

@@ -102,7 +102,7 @@ mod tests {
use rand_distr::Uniform;
use crate::ring::matrix_vec_product;
use crate::ring::PR;
use crate::ring::Rq;
#[test]
fn roots_of_unity() -> Result<()> {
@@ -136,7 +136,7 @@ mod tests {
let mut rng = rand::thread_rng();
let uniform_distr = Uniform::new(0_f64, Q as f64);
let a = PR::<Q, N>::rand_f64(&mut rng, uniform_distr)?;
let a = Rq::<Q, N>::rand_f64(&mut rng, uniform_distr)?;
// let a = PR::<Q, N>::new_from_u64(vec![36, 21, 9, 19]);
// let a_padded_coeffs: [Zq<Q>; 2 * N] =
@@ -148,7 +148,7 @@ mod tests {
let a_intt: Vec<Zq<Q>> = matrix_vec_product(&v_inv, &a_ntt)?;
assert_eq!(a_intt, a_padded);
let a_intt_arr: [Zq<Q>; N] = std::array::from_fn(|i| a_intt[i]);
assert_eq!(PR::new(a_intt_arr, None), a);
assert_eq!(Rq::new(a_intt_arr, None), a);
Ok(())
}
@@ -162,7 +162,7 @@ mod tests {
let a: Vec<Zq<Q>> = vec![256, 256, 256, 256, 0, 0, 0, 0]
.iter()
.map(|&e| Zq::new(e))
.map(|&e| Zq::from_u64(e))
.collect();
let a_ntt = matrix_vec_product(&ntt.ntt, &a)?;
let a_intt = matrix_vec_product(&ntt.intt, &a_ntt)?;
@@ -181,7 +181,7 @@ mod tests {
let ntt = NTT::<Q, N>::new()?;
let rng = rand::thread_rng();
let a = PR::<Q, { 2 * N }>::rand_f64(rng, Uniform::new(0_f64, (Q - 1) as f64))?;
let a = Rq::<Q, { 2 * N }>::rand_f64(rng, Uniform::new(0_f64, (Q - 1) as f64))?;
let a = a.coeffs;
dbg!(&a);
let a_ntt = matrix_vec_product(&ntt.ntt, &a.to_vec())?;

17
arithmetic/src/ntt.rs
View File

@@ -115,19 +115,20 @@ const fn roots_of_unity_inv<const Q: u64, const N: usize>(v: [Zq<Q>; N]) -> [Zq<
/// returns x^k mod Q
const fn const_exp_mod<const Q: u64>(x: u64, k: u64) -> u64 {
let mut r = 1u64;
let mut x = x;
let mut k = k;
x = x % Q;
// work on u128 to avoid overflow
let mut r = 1u128;
let mut x = x as u128;
let mut k = k as u128;
x = x % Q as u128;
// exponentiation by square strategy
while k > 0 {
if k % 2 == 1 {
r = (r * x) % Q;
r = (r * x) % Q as u128;
}
x = (x * x) % Q;
x = (x * x) % Q as u128;
k /= 2;
}
r
r as u64
}
/// returns x^-1 mod Q
@@ -149,7 +150,7 @@ mod tests {
const N: usize = 4;
let a: [u64; N] = [1u64, 2, 3, 4];
let a: [Zq<Q>; N] = array::from_fn(|i| Zq::new(a[i]));
let a: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(a[i]));
let a_ntt = NTT::<Q, N>::ntt(a);

482
arithmetic/src/ring.rs
View File

@@ -1,3 +1,7 @@
//! Polynomial ring Z[X]/(X^N+1)
//!
use anyhow::{anyhow, Result};
use rand::{distributions::Distribution, Rng};
use std::array;
use std::fmt;
@@ -5,432 +9,176 @@ use std::ops;
use crate::ntt::NTT;
use crate::zq::Zq;
use anyhow::{anyhow, Result};
// PolynomialRing element, where the PolynomialRing is R = Z_q[X]/(X^n +1)
#[derive(Clone, Copy)]
pub struct PR<const Q: u64, const N: usize> {
pub(crate) coeffs: [Zq<Q>; N],
// PolynomialRing element, where the PolynomialRing is R = Z[X]/(X^n +1)
#[derive(Clone, Copy, Debug)]
pub struct R<const N: usize>([i64; N]);
// evals are set when doig a PRxPR multiplication, so it can be reused in future
// multiplications avoiding recomputing it
pub(crate) evals: Option<[Zq<Q>; N]>,
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())
}
}
// 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
impl<const N: usize> R<N> {
pub fn coeffs(&self) -> [i64; N] {
self.0
}
pub fn to_rq<const Q: u64>(self) -> crate::Rq<Q, N> {
crate::Rq::<Q, N>::from(self)
}
pub fn from_vec(coeffs: Vec<i64>) -> Self {
let mut p = coeffs;
modulus::<N>(&mut p);
Self(array::from_fn(|i| p[i]))
}
// this method is mostly for tests
pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
let coeffs_i64 = coeffs.iter().map(|c| *c as i64).collect();
Self::from_vec(coeffs_i64)
}
pub fn from_vec_f64(coeffs: Vec<f64>) -> Self {
let coeffs_i64 = coeffs.iter().map(|c| c.round() as i64).collect();
Self::from_vec(coeffs_i64)
}
pub fn new(coeffs: [i64; N]) -> Self {
Self(coeffs)
}
pub fn mul_by_i64(&self, s: i64) -> Self {
Self(array::from_fn(|i| self.0[i] * s))
}
// performs the multiplication and division over f64, and then it rounds the
// result, only applying the mod Q at the end
pub fn mul_div_round<const Q: u64>(&self, num: u64, den: u64) -> crate::Rq<Q, N> {
let r: Vec<f64> = self
.coeffs()
.iter()
.map(|e| ((num as f64 * *e as f64) / den as f64).round())
.collect();
crate::Rq::<Q, N>::from_vec_f64(r)
}
}
pub fn mul_div_round<const Q: u64, const N: usize>(
v: Vec<i64>,
num: u64,
den: u64,
) -> crate::Rq<Q, N> {
// dbg!(&v);
let r: Vec<f64> = v
.iter()
.map(|e| ((num as f64 * *e as f64) / den as f64).round())
.collect();
// dbg!(&r);
crate::Rq::<Q, N>::from_vec_f64(r)
}
// apply mod (X^N+1)
pub fn modulus<const Q: u64, const N: usize>(p: &mut Vec<Zq<Q>>) {
pub fn modulus<const N: usize>(p: &mut Vec<i64>) {
if p.len() < N {
return;
}
for i in N..p.len() {
p[i - N] = p[i - N].clone() - p[i].clone();
p[i] = Zq(0);
p[i] = 0;
}
p.truncate(N);
}
// PR stands for PolynomialRing
impl<const Q: u64, const N: usize> PR<Q, N> {
pub fn coeffs(&self) -> [Zq<Q>; N] {
self.coeffs
}
pub fn from_vec(coeffs: Vec<Zq<Q>>) -> Self {
let mut p = coeffs;
modulus::<Q, N>(&mut p);
let coeffs = array::from_fn(|i| p[i]);
Self {
coeffs,
evals: None,
}
}
// this method is mostly for tests
pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
let coeffs_mod_q = coeffs.iter().map(|c| Zq::new(*c)).collect();
Self::from_vec(coeffs_mod_q)
}
pub fn new(coeffs: [Zq<Q>; N], evals: Option<[Zq<Q>; N]>) -> Self {
Self { coeffs, evals }
}
pub fn rand_abs(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
Ok(Self {
coeffs,
evals: None,
})
}
pub fn rand_f64(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng)));
Ok(Self {
coeffs,
evals: None,
})
}
pub fn rand_u64(mut rng: impl Rng, dist: impl Distribution<u64>) -> Result<Self> {
let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::new(dist.sample(&mut rng)));
Ok(Self {
coeffs,
evals: None,
})
}
// WIP. returns random v \in {0,1}. // TODO {-1, 0, 1}
pub fn rand_bin(mut rng: impl Rng, dist: impl Distribution<bool>) -> Result<Self> {
let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_bool(dist.sample(&mut rng)));
Ok(PR {
coeffs,
evals: None,
})
}
// Warning: this method assumes Q < P
pub fn remodule<const P: u64>(&self) -> PR<P, N> {
assert!(Q < P);
PR::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
}
// TODO review if needed, or if with this interface
pub fn mul_by_matrix(&self, m: &Vec<Vec<Zq<Q>>>) -> Result<Vec<Zq<Q>>> {
matrix_vec_product(m, &self.coeffs.to_vec())
}
pub fn mul_by_zq(&self, s: &Zq<Q>) -> Self {
Self {
coeffs: array::from_fn(|i| self.coeffs[i] * *s),
evals: None,
}
}
pub fn mul_by_u64(&self, s: u64) -> Self {
let s = Zq::new(s);
Self {
coeffs: array::from_fn(|i| self.coeffs[i] * s),
// coeffs: self.coeffs.iter().map(|&e| e * s).collect(),
evals: None,
}
}
pub fn mul_by_f64(&self, s: f64) -> Self {
Self {
coeffs: array::from_fn(|i| Zq::from_f64(self.coeffs[i].0 as f64 * s)),
evals: None,
}
}
pub fn mul(&mut self, rhs: &mut Self) -> Self {
mul_mut(self, rhs)
}
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
// TODO simplify
let mut str = "";
let mut zero = true;
for (i, coeff) in self.coeffs.iter().enumerate().rev() {
if coeff.0 == 0 {
continue;
}
zero = false;
f.write_str(str)?;
if coeff.0 != 1 {
f.write_str(coeff.0.to_string().as_str())?;
if i > 0 {
f.write_str("*")?;
}
}
if coeff.0 == 1 && i == 0 {
f.write_str(coeff.0.to_string().as_str())?;
}
if i == 1 {
f.write_str("x")?;
} else if i > 1 {
f.write_str("x^")?;
f.write_str(i.to_string().as_str())?;
}
str = " + ";
}
if zero {
f.write_str("0")?;
}
f.write_str(" mod Z_")?;
f.write_str(Q.to_string().as_str())?;
f.write_str("/(X^")?;
f.write_str(N.to_string().as_str())?;
f.write_str("+1)")?;
Ok(())
}
}
pub fn matrix_vec_product<const Q: u64>(m: &Vec<Vec<Zq<Q>>>, v: &Vec<Zq<Q>>) -> Result<Vec<Zq<Q>>> {
// assert_eq!(m.len(), m[0].len()); // TODO change to returning err
// assert_eq!(m.len(), v.len());
if m.len() != m[0].len() {
return Err(anyhow!("expected 'm' to be a square matrix"));
}
if m.len() != v.len() {
return Err(anyhow!(
"m.len: {} should be equal to v.len(): {}",
m.len(),
v.len(),
));
}
Ok(m.iter()
.map(|row| {
row.iter()
.zip(v.iter())
.map(|(&row_i, &v_i)| row_i * v_i)
.sum()
})
.collect::<Vec<Zq<Q>>>())
}
pub fn transpose<const Q: u64>(m: &[Vec<Zq<Q>>]) -> Vec<Vec<Zq<Q>>> {
// TODO case when m[0].len()=0
// TODO non square matrix
let mut r: Vec<Vec<Zq<Q>>> = vec![vec![Zq(0); m[0].len()]; m.len()];
for (i, m_row) in m.iter().enumerate() {
for (j, m_ij) in m_row.iter().enumerate() {
r[j][i] = *m_ij;
}
}
r
}
impl<const Q: u64, const N: usize> PartialEq for PR<Q, N> {
impl<const N: usize> PartialEq for R<N> {
fn eq(&self, other: &Self) -> bool {
self.coeffs == other.coeffs
self.0 == other.0
}
}
impl<const Q: u64, const N: usize> ops::Add<PR<Q, N>> for PR<Q, N> {
impl<const N: usize> ops::Add<R<N>> for R<N> {
type Output = Self;
fn add(self, rhs: Self) -> Self {
Self {
coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
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 +
Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
}
}
impl<const Q: u64, const N: usize> ops::Add<&PR<Q, N>> for &PR<Q, N> {
type Output = PR<Q, N>;
impl<const N: usize> ops::Add<&R<N>> for &R<N> {
type Output = R<N>;
fn add(self, rhs: &PR<Q, N>) -> Self::Output {
PR {
coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
evals: None,
}
fn add(self, rhs: &R<N>) -> Self::Output {
R(array::from_fn(|i| self.0[i] + rhs.0[i]))
}
}
impl<const Q: u64, const N: usize> ops::Sub<PR<Q, N>> for PR<Q, N> {
impl<const N: usize> ops::Sub<R<N>> for R<N> {
type Output = Self;
fn sub(self, rhs: Self) -> Self {
Self {
coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
evals: None,
}
Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
}
}
impl<const Q: u64, const N: usize> ops::Sub<&PR<Q, N>> for &PR<Q, N> {
type Output = PR<Q, N>;
impl<const N: usize> ops::Sub<&R<N>> for &R<N> {
type Output = R<N>;
fn sub(self, rhs: &PR<Q, N>) -> Self::Output {
PR {
coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
evals: None,
}
fn sub(self, rhs: &R<N>) -> Self::Output {
R(array::from_fn(|i| self.0[i] - rhs.0[i]))
}
}
impl<const Q: u64, const N: usize> ops::Mul<PR<Q, N>> for PR<Q, N> {
impl<const N: usize> ops::Mul<R<N>> for R<N> {
type Output = Self;
fn mul(self, rhs: Self) -> Self {
mul(&self, &rhs)
naive_poly_mul(&self, &rhs)
}
}
impl<const Q: u64, const N: usize> ops::Mul<&PR<Q, N>> for &PR<Q, N> {
type Output = PR<Q, N>;
impl<const N: usize> ops::Mul<&R<N>> for &R<N> {
type Output = R<N>;
fn mul(self, rhs: &PR<Q, N>) -> Self::Output {
mul(self, rhs)
fn mul(self, rhs: &R<N>) -> Self::Output {
naive_poly_mul(self, rhs)
}
}
// mul by Zq element
impl<const Q: u64, const N: usize> ops::Mul<Zq<Q>> for PR<Q, N> {
type Output = Self;
fn mul(self, s: Zq<Q>) -> Self {
self.mul_by_zq(&s)
// TODO with NTT(?)
pub fn naive_poly_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> R<N> {
let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
let mut result: Vec<i128> = vec![0; (N * 2) - 1];
for i in 0..N {
for j in 0..N {
result[i + j] = result[i + j] + poly1[i] * poly2[j];
}
}
}
impl<const Q: u64, const N: usize> ops::Mul<&Zq<Q>> for &PR<Q, N> {
type Output = PR<Q, N>;
fn mul(self, s: &Zq<Q>) -> Self::Output {
self.mul_by_zq(s)
}
// apply mod (X^N + 1))
R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
}
pub fn naive_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> Vec<i64> {
let poly1: Vec<i128> = poly1.0.iter().map(|c| *c as i128).collect();
let poly2: Vec<i128> = poly2.0.iter().map(|c| *c as i128).collect();
let mut result = vec![0; (N * 2) - 1];
for i in 0..N {
for j in 0..N {
result[i + j] = result[i + j] + poly1[i] * poly2[j];
}
}
result.iter().map(|c| *c as i64).collect()
}
// mul by u64
impl<const Q: u64, const N: usize> ops::Mul<u64> for PR<Q, N> {
impl<const N: usize> ops::Mul<u64> for R<N> {
type Output = Self;
fn mul(self, s: u64) -> Self {
self.mul_by_u64(s)
self.mul_by_i64(s as i64)
}
}
impl<const Q: u64, const N: usize> ops::Mul<&u64> for &PR<Q, N> {
type Output = PR<Q, N>;
impl<const N: usize> ops::Mul<&u64> for &R<N> {
type Output = R<N>;
fn mul(self, s: &u64) -> Self::Output {
self.mul_by_u64(*s)
self.mul_by_i64(*s as i64)
}
}
impl<const Q: u64, const N: usize> ops::Neg for PR<Q, N> {
impl<const N: usize> ops::Neg for R<N> {
type Output = Self;
fn neg(self) -> Self::Output {
Self {
coeffs: array::from_fn(|i| -self.coeffs[i]),
evals: None,
}
}
}
fn mul_mut<const Q: u64, const N: usize>(lhs: &mut PR<Q, N>, rhs: &mut PR<Q, N>) -> PR<Q, N> {
// reuse evaluations if already computed
if !lhs.evals.is_some() {
lhs.evals = Some(NTT::<Q, N>::ntt(lhs.coeffs));
};
if !rhs.evals.is_some() {
rhs.evals = Some(NTT::<Q, N>::ntt(rhs.coeffs));
};
let lhs_evals = lhs.evals.unwrap();
let rhs_evals = rhs.evals.unwrap();
let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
let c = NTT::<Q, { N }>::intt(c_ntt);
PR::new(c, Some(c_ntt))
}
fn mul<const Q: u64, const N: usize>(lhs: &PR<Q, N>, rhs: &PR<Q, N>) -> PR<Q, N> {
// reuse evaluations if already computed
let lhs_evals = if lhs.evals.is_some() {
lhs.evals.unwrap()
} else {
NTT::<Q, N>::ntt(lhs.coeffs)
};
let rhs_evals = if rhs.evals.is_some() {
rhs.evals.unwrap()
} else {
NTT::<Q, N>::ntt(rhs.coeffs)
};
let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
let c = NTT::<Q, { N }>::intt(c_ntt);
PR::new(c, Some(c_ntt))
}
impl<const Q: u64, const N: usize> fmt::Display for PR<Q, N> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.fmt(f)?;
Ok(())
}
}
impl<const Q: u64, const N: usize> fmt::Debug for PR<Q, N> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.fmt(f)?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn poly_ring() {
// the test values used are generated with SageMath
const Q: u64 = 7;
const N: usize = 3;
// p = 1x + 2x^2 + 3x^3 + 4 x^4 + 5 x^5 in R=Z_q[X]/(X^n +1)
let p = PR::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
// try with coefficients bigger than Q
let p = PR::<Q, N>::from_vec_u64(vec![0u64, 1, Q + 2, 3, 4, 5]);
assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
// try with other ring
let p = PR::<7, 4>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
assert_eq!(p.to_string(), "3*x^3 + 2*x^2 + 3*x + 3 mod Z_7/(X^4+1)");
let p = PR::<Q, N>::from_vec_u64(vec![0u64, 0, 0, 0, 4, 5]);
assert_eq!(p.to_string(), "2*x^2 + 3*x mod Z_7/(X^3+1)");
let p = PR::<Q, N>::from_vec_u64(vec![5u64, 4, 5, 2, 1, 0]);
assert_eq!(p.to_string(), "5*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
let a = PR::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
assert_eq!(a.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
let b = PR::<Q, N>::from_vec_u64(vec![5u64, 4, 3, 2, 1, 0]);
assert_eq!(b.to_string(), "3*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
// add
assert_eq!((a.clone() + b.clone()).to_string(), "0 mod Z_7/(X^3+1)");
assert_eq!((&a + &b).to_string(), "0 mod Z_7/(X^3+1)");
// assert_eq!((a.0.clone() + b.0.clone()).to_string(), "[0, 0, 0]"); // TODO
// sub
assert_eq!(
(a.clone() - b.clone()).to_string(),
"x^2 + x + 1 mod Z_7/(X^3+1)"
);
}
fn test_mul_opt<const Q: u64, const N: usize>(
a: [u64; N],
b: [u64; N],
expected_c: [u64; N],
) -> Result<()> {
let a: [Zq<Q>; N] = array::from_fn(|i| Zq::new(a[i]));
let mut a = PR::new(a, None);
let b: [Zq<Q>; N] = array::from_fn(|i| Zq::new(b[i]));
let mut b = PR::new(b, None);
let expected_c: [Zq<Q>; N] = array::from_fn(|i| Zq::new(expected_c[i]));
let expected_c = PR::new(expected_c, None);
let c = mul_mut(&mut a, &mut b);
assert_eq!(c, expected_c);
Ok(())
}
#[test]
fn test_mul() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 4;
let a: [u64; N] = [1u64, 2, 3, 4];
let b: [u64; N] = [1u64, 2, 3, 4];
let c: [u64; N] = [65513, 65517, 65531, 20];
test_mul_opt::<Q, N>(a, b, c)?;
let a: [u64; N] = [0u64, 0, 0, 2];
let b: [u64; N] = [0u64, 0, 0, 2];
let c: [u64; N] = [0u64, 0, 65533, 0];
test_mul_opt::<Q, N>(a, b, c)?;
// TODO more testvectors
Ok(())
Self(array::from_fn(|i| -self.0[i]))
}
}

504
arithmetic/src/ringq.rs Normal file
View File

@@ -0,0 +1,504 @@
//! Polynomial ring Z_q[X]/(X^N+1)
//!
use rand::{distributions::Distribution, Rng};
use std::array;
use std::fmt;
use std::ops;
use crate::ntt::NTT;
use crate::zq::{modulus_u64, Zq};
use anyhow::{anyhow, Result};
/// PolynomialRing element, where the PolynomialRing is R = Z_q[X]/(X^n +1)
/// The implementation assumes that q is prime.
#[derive(Clone, Copy)]
pub struct Rq<const Q: u64, const N: usize> {
pub(crate) coeffs: [Zq<Q>; N],
// evals are set when doig a PRxPR multiplication, so it can be reused in future
// multiplications avoiding recomputing it
pub(crate) evals: Option<[Zq<Q>; N]>,
}
// 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
impl<const Q: u64, const N: usize> From<crate::ring::R<N>> for Rq<Q, N> {
fn from(r: crate::ring::R<N>) -> Self {
Self::from_vec(
r.coeffs()
.iter()
.map(|e| Zq::<Q>::from_f64(*e as f64))
.collect(),
)
}
}
// apply mod (X^N+1)
pub fn modulus<const Q: u64, const N: usize>(p: &mut Vec<Zq<Q>>) {
if p.len() < N {
return;
}
for i in N..p.len() {
p[i - N] = p[i - N].clone() - p[i].clone();
p[i] = Zq(0);
}
p.truncate(N);
}
// PR stands for PolynomialRing
impl<const Q: u64, const N: usize> Rq<Q, N> {
pub fn coeffs(&self) -> [Zq<Q>; N] {
self.coeffs
}
pub fn to_r(self) -> crate::R<N> {
crate::R::<N>::from(self)
}
pub fn from_vec(coeffs: Vec<Zq<Q>>) -> Self {
let mut p = coeffs;
modulus::<Q, N>(&mut p);
let coeffs = array::from_fn(|i| p[i]);
Self {
coeffs,
evals: None,
}
}
// this method is mostly for tests
pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
let coeffs_mod_q = coeffs.iter().map(|c| Zq::from_u64(*c)).collect();
Self::from_vec(coeffs_mod_q)
}
pub fn from_vec_f64(coeffs: Vec<f64>) -> Self {
let coeffs_mod_q = coeffs.iter().map(|c| Zq::from_f64(*c)).collect();
Self::from_vec(coeffs_mod_q)
}
pub fn from_vec_i64(coeffs: Vec<i64>) -> Self {
let coeffs_mod_q = coeffs.iter().map(|c| Zq::from_f64(*c as f64)).collect();
Self::from_vec(coeffs_mod_q)
}
pub fn new(coeffs: [Zq<Q>; N], evals: Option<[Zq<Q>; N]>) -> Self {
Self { coeffs, evals }
}
pub fn rand_abs(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
Ok(Self {
coeffs,
evals: None,
})
}
pub fn rand_f64_abs(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng).abs()));
Ok(Self {
coeffs,
evals: None,
})
}
pub fn rand_f64(mut rng: impl Rng, dist: impl Distribution<f64>) -> Result<Self> {
let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng)));
Ok(Self {
coeffs,
evals: None,
})
}
pub fn rand_u64(mut rng: impl Rng, dist: impl Distribution<u64>) -> Result<Self> {
let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_u64(dist.sample(&mut rng)));
Ok(Self {
coeffs,
evals: None,
})
}
// WIP. returns random v \in {0,1}. // TODO {-1, 0, 1}
pub fn rand_bin(mut rng: impl Rng, dist: impl Distribution<bool>) -> Result<Self> {
let coeffs: [Zq<Q>; N] = array::from_fn(|_| Zq::from_bool(dist.sample(&mut rng)));
Ok(Rq {
coeffs,
evals: None,
})
}
// Warning: this method will behave differently depending on the values P and Q:
// if Q<P, it just 'renames' the modulus parameter to P
// if Q>=P, it crops to mod P
pub fn remodule<const P: u64>(&self) -> Rq<P, N> {
Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
}
// applies mod(T) to all coefficients of self
pub fn coeffs_mod<const T: u64>(&self) -> Self {
Rq::<Q, N>::from_vec_u64(
self.coeffs()
.iter()
.map(|m_i| modulus_u64::<T>(m_i.0))
.collect(),
)
}
// TODO review if needed, or if with this interface
pub fn mul_by_matrix(&self, m: &Vec<Vec<Zq<Q>>>) -> Result<Vec<Zq<Q>>> {
matrix_vec_product(m, &self.coeffs.to_vec())
}
pub fn mul_by_zq(&self, s: &Zq<Q>) -> Self {
Self {
coeffs: array::from_fn(|i| self.coeffs[i] * *s),
evals: None,
}
}
pub fn mul_by_u64(&self, s: u64) -> Self {
let s = Zq::from_u64(s);
Self {
coeffs: array::from_fn(|i| self.coeffs[i] * s),
// coeffs: self.coeffs.iter().map(|&e| e * s).collect(),
evals: None,
}
}
pub fn mul_by_f64(&self, s: f64) -> Self {
Self {
coeffs: array::from_fn(|i| Zq::from_f64(self.coeffs[i].0 as f64 * s)),
evals: None,
}
}
pub fn mul(&mut self, rhs: &mut Self) -> Self {
mul_mut(self, rhs)
}
// divides by the given scalar 's' and rounds, returning a Rq<Q,N>
// TODO rm
pub fn div_round(&self, s: u64) -> Self {
let r: Vec<f64> = self
.coeffs()
.iter()
.map(|e| (e.0 as f64 / s as f64).round())
.collect();
Rq::<Q, N>::from_vec_f64(r)
}
// returns [ [(num/den) * self].round() ] mod q
// ie. performs the multiplication and division over f64, and then it rounds the
// result, only applying the mod Q at the end
pub fn mul_div_round(&self, num: u64, den: u64) -> Self {
let r: Vec<f64> = self
.coeffs()
.iter()
.map(|e| ((num as f64 * e.0 as f64) / den as f64).round())
.collect();
Rq::<Q, N>::from_vec_f64(r)
}
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
// TODO simplify
let mut str = "";
let mut zero = true;
for (i, coeff) in self.coeffs.iter().enumerate().rev() {
if coeff.0 == 0 {
continue;
}
zero = false;
f.write_str(str)?;
if coeff.0 != 1 {
f.write_str(coeff.0.to_string().as_str())?;
if i > 0 {
f.write_str("*")?;
}
}
if coeff.0 == 1 && i == 0 {
f.write_str(coeff.0.to_string().as_str())?;
}
if i == 1 {
f.write_str("x")?;
} else if i > 1 {
f.write_str("x^")?;
f.write_str(i.to_string().as_str())?;
}
str = " + ";
}
if zero {
f.write_str("0")?;
}
f.write_str(" mod Z_")?;
f.write_str(Q.to_string().as_str())?;
f.write_str("/(X^")?;
f.write_str(N.to_string().as_str())?;
f.write_str("+1)")?;
Ok(())
}
pub fn infinity_norm(&self) -> u64 {
self.coeffs().iter().map(|x| x.0).fold(0, |a, b| a.max(b))
}
}
pub fn matrix_vec_product<const Q: u64>(m: &Vec<Vec<Zq<Q>>>, v: &Vec<Zq<Q>>) -> Result<Vec<Zq<Q>>> {
// assert_eq!(m.len(), m[0].len()); // TODO change to returning err
// assert_eq!(m.len(), v.len());
if m.len() != m[0].len() {
return Err(anyhow!("expected 'm' to be a square matrix"));
}
if m.len() != v.len() {
return Err(anyhow!(
"m.len: {} should be equal to v.len(): {}",
m.len(),
v.len(),
));
}
Ok(m.iter()
.map(|row| {
row.iter()
.zip(v.iter())
.map(|(&row_i, &v_i)| row_i * v_i)
.sum()
})
.collect::<Vec<Zq<Q>>>())
}
pub fn transpose<const Q: u64>(m: &[Vec<Zq<Q>>]) -> Vec<Vec<Zq<Q>>> {
// TODO case when m[0].len()=0
// TODO non square matrix
let mut r: Vec<Vec<Zq<Q>>> = vec![vec![Zq(0); m[0].len()]; m.len()];
for (i, m_row) in m.iter().enumerate() {
for (j, m_ij) in m_row.iter().enumerate() {
r[j][i] = *m_ij;
}
}
r
}
impl<const Q: u64, const N: usize> PartialEq for Rq<Q, N> {
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> {
type Output = Self;
fn add(self, rhs: Self) -> Self {
Self {
coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
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<const Q: u64, const N: usize> ops::Add<&Rq<Q, N>> for &Rq<Q, N> {
type Output = Rq<Q, N>;
fn add(self, rhs: &Rq<Q, N>) -> Self::Output {
Rq {
coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
evals: None,
}
}
}
impl<const Q: u64, const N: usize> ops::Sub<Rq<Q, N>> for Rq<Q, N> {
type Output = Self;
fn sub(self, rhs: Self) -> Self {
Self {
coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
evals: None,
}
}
}
impl<const Q: u64, const N: usize> ops::Sub<&Rq<Q, N>> for &Rq<Q, N> {
type Output = Rq<Q, N>;
fn sub(self, rhs: &Rq<Q, N>) -> Self::Output {
Rq {
coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
evals: None,
}
}
}
impl<const Q: u64, const N: usize> ops::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> {
type Output = Rq<Q, N>;
fn mul(self, rhs: &Rq<Q, N>) -> Self::Output {
mul(self, rhs)
}
}
// mul by Zq element
impl<const Q: u64, const N: usize> ops::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> {
type Output = Rq<Q, N>;
fn mul(self, s: &Zq<Q>) -> Self::Output {
self.mul_by_zq(s)
}
}
// mul by u64
impl<const Q: u64, const N: usize> ops::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> {
type Output = Rq<Q, N>;
fn mul(self, s: &u64) -> Self::Output {
self.mul_by_u64(*s)
}
}
impl<const Q: u64, const N: usize> ops::Neg for Rq<Q, N> {
type Output = Self;
fn neg(self) -> Self::Output {
Self {
coeffs: array::from_fn(|i| -self.coeffs[i]),
evals: None,
}
}
}
fn mul_mut<const Q: u64, const N: usize>(lhs: &mut Rq<Q, N>, rhs: &mut Rq<Q, N>) -> Rq<Q, N> {
// reuse evaluations if already computed
if !lhs.evals.is_some() {
lhs.evals = Some(NTT::<Q, N>::ntt(lhs.coeffs));
};
if !rhs.evals.is_some() {
rhs.evals = Some(NTT::<Q, N>::ntt(rhs.coeffs));
};
let lhs_evals = lhs.evals.unwrap();
let rhs_evals = rhs.evals.unwrap();
let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
let c = NTT::<Q, { N }>::intt(c_ntt);
Rq::new(c, Some(c_ntt))
}
fn mul<const Q: u64, const N: usize>(lhs: &Rq<Q, N>, rhs: &Rq<Q, N>) -> Rq<Q, N> {
// reuse evaluations if already computed
let lhs_evals = if lhs.evals.is_some() {
lhs.evals.unwrap()
} else {
NTT::<Q, N>::ntt(lhs.coeffs)
};
let rhs_evals = if rhs.evals.is_some() {
rhs.evals.unwrap()
} else {
NTT::<Q, N>::ntt(rhs.coeffs)
};
let c_ntt: [Zq<Q>; N] = array::from_fn(|i| lhs_evals[i] * rhs_evals[i]);
let c = NTT::<Q, { N }>::intt(c_ntt);
Rq::new(c, Some(c_ntt))
}
impl<const Q: u64, const N: usize> fmt::Display for Rq<Q, N> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.fmt(f)?;
Ok(())
}
}
impl<const Q: u64, const N: usize> fmt::Debug for Rq<Q, N> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.fmt(f)?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn poly_ring() {
// the test values used are generated with SageMath
const Q: u64 = 7;
const N: usize = 3;
// p = 1x + 2x^2 + 3x^3 + 4 x^4 + 5 x^5 in R=Z_q[X]/(X^n +1)
let p = Rq::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
// try with coefficients bigger than Q
let p = Rq::<Q, N>::from_vec_u64(vec![0u64, 1, Q + 2, 3, 4, 5]);
assert_eq!(p.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
// try with other ring
let p = Rq::<7, 4>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
assert_eq!(p.to_string(), "3*x^3 + 2*x^2 + 3*x + 3 mod Z_7/(X^4+1)");
let p = Rq::<Q, N>::from_vec_u64(vec![0u64, 0, 0, 0, 4, 5]);
assert_eq!(p.to_string(), "2*x^2 + 3*x mod Z_7/(X^3+1)");
let p = Rq::<Q, N>::from_vec_u64(vec![5u64, 4, 5, 2, 1, 0]);
assert_eq!(p.to_string(), "5*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
let a = Rq::<Q, N>::from_vec_u64(vec![0u64, 1, 2, 3, 4, 5]);
assert_eq!(a.to_string(), "4*x^2 + 4*x + 4 mod Z_7/(X^3+1)");
let b = Rq::<Q, N>::from_vec_u64(vec![5u64, 4, 3, 2, 1, 0]);
assert_eq!(b.to_string(), "3*x^2 + 3*x + 3 mod Z_7/(X^3+1)");
// add
assert_eq!((a.clone() + b.clone()).to_string(), "0 mod Z_7/(X^3+1)");
assert_eq!((&a + &b).to_string(), "0 mod Z_7/(X^3+1)");
// assert_eq!((a.0.clone() + b.0.clone()).to_string(), "[0, 0, 0]"); // TODO
// sub
assert_eq!(
(a.clone() - b.clone()).to_string(),
"x^2 + x + 1 mod Z_7/(X^3+1)"
);
}
fn test_mul_opt<const Q: u64, const N: usize>(
a: [u64; N],
b: [u64; N],
expected_c: [u64; N],
) -> Result<()> {
let a: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(a[i]));
let mut a = Rq::new(a, None);
let b: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(b[i]));
let mut b = Rq::new(b, None);
let expected_c: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(expected_c[i]));
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_mul() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 4;
let a: [u64; N] = [1u64, 2, 3, 4];
let b: [u64; N] = [1u64, 2, 3, 4];
let c: [u64; N] = [65513, 65517, 65531, 20];
test_mul_opt::<Q, N>(a, b, c)?;
let a: [u64; N] = [0u64, 0, 0, 2];
let b: [u64; N] = [0u64, 0, 0, 2];
let c: [u64; N] = [0u64, 0, 65533, 0];
test_mul_opt::<Q, N>(a, b, c)?;
// TODO more testvectors
Ok(())
}
}

28
arithmetic/src/zq.rs
View File

@@ -12,22 +12,36 @@ pub struct Zq<const Q: u64>(pub u64);
// }
// }
pub(crate) fn modulus_u64<const Q: u64>(e: u64) -> u64 {
(e % Q + Q) % Q
}
impl<const Q: u64> Zq<Q> {
pub fn new(e: u64) -> Self {
pub fn from_u64(e: u64) -> Self {
if e >= Q {
return Zq(e % Q);
// (e % Q + Q) % Q
return Zq(modulus_u64::<Q>(e));
// return Zq(e % Q);
}
Zq(e)
}
pub fn from_f64(e: f64) -> Self {
// WIP method
let e: i64 = e.round() as i64;
if e < 0 {
return Zq((Q as i64 + e) as u64);
} else if e >= Q as i64 {
return Zq((e % Q as i64) as u64);
let q = Q as i64;
if e < 0 || e >= q {
return Zq(((e % q + q) % q) as u64);
}
Zq(e as u64)
// if e < 0 {
// // dbg!(&e);
// // dbg!(Zq::<Q>(((Q as i64 + e) % Q as i64) as u64));
// // return Zq(((Q as i64 + e) % Q as i64) as u64);
// return Zq(e as u64 % Q);
// } else if e >= Q as i64 {
// return Zq((e % Q as i64) as u64);
// }
// Zq(e as u64)
}
pub fn from_bool(b: bool) -> Self {
if b {
@@ -83,7 +97,7 @@ impl<const Q: u64> Zq<Q> {
// if t < 0 {
// t = t + q;
// }
return Zq::new(t);
return Zq::from_u64(t);
}
pub fn inv(self) -> Zq<Q> {
let (g, x, _) = Self::egcd(self.0 as i128, Q as i128);