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

1
Cargo.toml
View File

@@ -10,3 +10,4 @@ resolver = "2"
anyhow = "1.0.56" anyhow = "1.0.56"
rand = "0.8.5" rand = "0.8.5"
rand_distr = "0.4.3" rand_distr = "0.4.3"
itertools = "0.14.0"

2
README.md
View File

@@ -1,4 +1,4 @@
# fhe-study # fhe-study
Code done while studying some FHE papers. Code done while studying some FHE papers, with the idea of doing implementations from scratch.
- arithmetic: contains $\mathbb{Z}_q$ and $\mathbb{Z}_q[X]/(X^N+1)$ arithmetic implementations, together with the NTT implementation. - arithmetic: contains $\mathbb{Z}_q$ and $\mathbb{Z}_q[X]/(X^N+1)$ arithmetic implementations, together with the NTT implementation.

6
arithmetic/src/lib.rs
View File

@@ -4,11 +4,13 @@
#![allow(clippy::upper_case_acronyms)] #![allow(clippy::upper_case_acronyms)]
#![allow(dead_code)] // TMP #![allow(dead_code)] // TMP
mod naive; // TODO rm mod naive_ntt; // TODO rm
pub mod ntt; pub mod ntt;
pub mod ring; pub mod ring;
pub mod ringq;
pub mod zq; pub mod zq;
pub use ntt::NTT; pub use ntt::NTT;
pub use ring::PR; pub use ring::R;
pub use ringq::Rq;
pub use zq::Zq; 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 rand_distr::Uniform;
use crate::ring::matrix_vec_product; use crate::ring::matrix_vec_product;
use crate::ring::PR; use crate::ring::Rq;
#[test] #[test]
fn roots_of_unity() -> Result<()> { fn roots_of_unity() -> Result<()> {
@@ -136,7 +136,7 @@ mod tests {
let mut rng = rand::thread_rng(); let mut rng = rand::thread_rng();
let uniform_distr = Uniform::new(0_f64, Q as f64); 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 = PR::<Q, N>::new_from_u64(vec![36, 21, 9, 19]);
// let a_padded_coeffs: [Zq<Q>; 2 * N] = // 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)?; let a_intt: Vec<Zq<Q>> = matrix_vec_product(&v_inv, &a_ntt)?;
assert_eq!(a_intt, a_padded); assert_eq!(a_intt, a_padded);
let a_intt_arr: [Zq<Q>; N] = std::array::from_fn(|i| a_intt[i]); 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(()) Ok(())
} }
@@ -162,7 +162,7 @@ mod tests {
let a: Vec<Zq<Q>> = vec![256, 256, 256, 256, 0, 0, 0, 0] let a: Vec<Zq<Q>> = vec![256, 256, 256, 256, 0, 0, 0, 0]
.iter() .iter()
.map(|&e| Zq::new(e)) .map(|&e| Zq::from_u64(e))
.collect(); .collect();
let a_ntt = matrix_vec_product(&ntt.ntt, &a)?; let a_ntt = matrix_vec_product(&ntt.ntt, &a)?;
let a_intt = matrix_vec_product(&ntt.intt, &a_ntt)?; let a_intt = matrix_vec_product(&ntt.intt, &a_ntt)?;
@@ -181,7 +181,7 @@ mod tests {
let ntt = NTT::<Q, N>::new()?; let ntt = NTT::<Q, N>::new()?;
let rng = rand::thread_rng(); 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; let a = a.coeffs;
dbg!(&a); dbg!(&a);
let a_ntt = matrix_vec_product(&ntt.ntt, &a.to_vec())?; 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 /// returns x^k mod Q
const fn const_exp_mod<const Q: u64>(x: u64, k: u64) -> u64 { const fn const_exp_mod<const Q: u64>(x: u64, k: u64) -> u64 {
let mut r = 1u64; // work on u128 to avoid overflow
let mut x = x; let mut r = 1u128;
let mut k = k; let mut x = x as u128;
x = x % Q; let mut k = k as u128;
x = x % Q as u128;
// exponentiation by square strategy // exponentiation by square strategy
while k > 0 { while k > 0 {
if k % 2 == 1 { 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; k /= 2;
} }
r r as u64
} }
/// returns x^-1 mod Q /// returns x^-1 mod Q
@@ -149,7 +150,7 @@ mod tests {
const N: usize = 4; const N: usize = 4;
let a: [u64; N] = [1u64, 2, 3, 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); 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 rand::{distributions::Distribution, Rng};
use std::array; use std::array;
use std::fmt; use std::fmt;
@@ -5,432 +9,176 @@ use std::ops;
use crate::ntt::NTT; use crate::ntt::NTT;
use crate::zq::Zq; use crate::zq::Zq;
use anyhow::{anyhow, Result};
// PolynomialRing element, where the PolynomialRing is R = Z_q[X]/(X^n +1) // PolynomialRing element, where the PolynomialRing is R = Z[X]/(X^n +1)
#[derive(Clone, Copy)] #[derive(Clone, Copy, Debug)]
pub struct PR<const Q: u64, const N: usize> { pub struct R<const N: usize>([i64; N]);
pub(crate) coeffs: [Zq<Q>; N],
// evals are set when doig a PRxPR multiplication, so it can be reused in future impl<const Q: u64, const N: usize> From<crate::ringq::Rq<Q, N>> for R<N> {
// multiplications avoiding recomputing it fn from(rq: crate::ringq::Rq<Q, N>) -> Self {
pub(crate) evals: Option<[Zq<Q>; N]>, 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 impl<const N: usize> R<N> {
// use the trait and not need to add '<Q, N>' to their params 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) // 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 { if p.len() < N {
return; return;
} }
for i in N..p.len() { for i in N..p.len() {
p[i - N] = p[i - N].clone() - p[i].clone(); p[i - N] = p[i - N].clone() - p[i].clone();
p[i] = Zq(0); p[i] = 0;
} }
p.truncate(N); p.truncate(N);
} }
// PR stands for PolynomialRing impl<const N: usize> PartialEq for R<N> {
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> {
fn eq(&self, other: &Self) -> bool { 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; type Output = Self;
fn add(self, rhs: Self) -> Self { fn add(self, rhs: Self) -> Self {
Self { Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
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<&PR<Q, N>> for &PR<Q, N> { impl<const N: usize> ops::Add<&R<N>> for &R<N> {
type Output = PR<Q, N>; type Output = R<N>;
fn add(self, rhs: &PR<Q, N>) -> Self::Output { fn add(self, rhs: &R<N>) -> Self::Output {
PR { R(array::from_fn(|i| self.0[i] + rhs.0[i]))
coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
evals: None,
}
} }
} }
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; type Output = Self;
fn sub(self, rhs: Self) -> Self { fn sub(self, rhs: Self) -> Self {
Self { Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
evals: None,
}
} }
} }
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 = PR<Q, N>; type Output = R<N>;
fn sub(self, rhs: &PR<Q, N>) -> Self::Output { fn sub(self, rhs: &R<N>) -> Self::Output {
PR { R(array::from_fn(|i| self.0[i] - rhs.0[i]))
coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
evals: None,
}
} }
} }
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; type Output = Self;
fn mul(self, rhs: Self) -> 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> { impl<const N: usize> ops::Mul<&R<N>> for &R<N> {
type Output = PR<Q, N>; type Output = R<N>;
fn mul(self, rhs: &PR<Q, N>) -> Self::Output { fn mul(self, rhs: &R<N>) -> Self::Output {
mul(self, rhs) naive_poly_mul(self, rhs)
} }
} }
// mul by Zq element // TODO with NTT(?)
impl<const Q: u64, const N: usize> ops::Mul<Zq<Q>> for PR<Q, N> { pub fn naive_poly_mul<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> R<N> {
type Output = Self; 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();
fn mul(self, s: Zq<Q>) -> Self { let mut result: Vec<i128> = vec![0; (N * 2) - 1];
self.mul_by_zq(&s) 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 // 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; type Output = Self;
fn mul(self, s: u64) -> 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> { impl<const N: usize> ops::Mul<&u64> for &R<N> {
type Output = PR<Q, N>; type Output = R<N>;
fn mul(self, s: &u64) -> Self::Output { 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; type Output = Self;
fn neg(self) -> Self::Output { fn neg(self) -> Self::Output {
Self { Self(array::from_fn(|i| -self.0[i]))
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(())
} }
} }

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

1
bfv/Cargo.toml
View File

@@ -7,5 +7,6 @@ edition = "2024"
anyhow = { workspace = true } anyhow = { workspace = true }
rand = { workspace = true } rand = { workspace = true }
rand_distr = { workspace = true } rand_distr = { workspace = true }
itertools = { workspace = true }
arithmetic = { path="../arithmetic" } arithmetic = { path="../arithmetic" }

559
bfv/src/lib.rs
View File

@@ -10,26 +10,100 @@ use rand::Rng;
use rand_distr::{Normal, Uniform}; use rand_distr::{Normal, Uniform};
use std::ops; use std::ops;
use arithmetic::{Zq, PR}; use arithmetic::{Rq, Zq, R};
// error deviation for the Gaussian(Normal) distribution // error deviation for the Gaussian(Normal) distribution
// sigma=3.2 from: https://eprint.iacr.org/2022/162.pdf page 5 // sigma=3.2 from: https://eprint.iacr.org/2022/162.pdf page 5
const ERR_SIGMA: f64 = 3.2; const ERR_SIGMA: f64 = 3.2;
// const ERR_SIGMA: f64 = 0.0; // TODO WIP
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct SecretKey<const Q: u64, const N: usize>(PR<Q, N>); pub struct SecretKey<const Q: u64, const N: usize>(Rq<Q, N>);
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct PublicKey<const Q: u64, const N: usize>(PR<Q, N>, PR<Q, N>); pub struct PublicKey<const Q: u64, const N: usize>(Rq<Q, N>, Rq<Q, N>);
/// Relinearization key
#[derive(Clone, Debug)]
pub struct RLK<const PQ: u64, const N: usize>(Rq<PQ, N>, Rq<PQ, N>);
// impl<const PQ: u64, const N: usize> RLK<Q, PQ, N> {
// // const P: u64 = PQ / Q;
//
// // const PQ: u64 = P * Q;
// }
// RLWE ciphertext // RLWE ciphertext
#[derive(Clone, Debug)] #[derive(Clone, Debug)]
pub struct RLWE<const Q: u64, const N: usize>(PR<Q, N>, PR<Q, N>); pub struct RLWE<const Q: u64, const N: usize>(Rq<Q, N>, Rq<Q, N>);
impl<const Q: u64, const N: usize> RLWE<Q, N> { impl<const Q: u64, const N: usize> RLWE<Q, N> {
fn add(lhs: Self, rhs: Self) -> Self { fn add(lhs: Self, rhs: Self) -> Self {
RLWE::<Q, N>(lhs.0 + rhs.0, lhs.1 + rhs.1) RLWE::<Q, N>(lhs.0 + rhs.0, lhs.1 + rhs.1)
} }
pub fn remodule<const P: u64>(&self) -> RLWE<P, N> {
let x = self.0.remodule::<P>();
let y = self.1.remodule::<P>();
RLWE::<P, N>(x, y)
}
fn tensor<const PQ: u64, const T: u64>(a: &Self, b: &Self) -> (Rq<Q, N>, Rq<Q, N>, Rq<Q, N>) {
// expand Q->PQ // TODO rm
// get the coefficients in Z, ie. interpret a,b \in R (instead of R_q)
let a0: R<N> = a.0.to_r();
let a1: R<N> = a.1.to_r();
let b0: R<N> = b.0.to_r();
let b1: R<N> = b.1.to_r();
// tensor (\in R)
use arithmetic::ring::naive_mul;
let c0: Vec<i64> = naive_mul(&a0, &b0);
let c1_l: Vec<i64> = naive_mul(&a0, &b1);
let c1_r = naive_mul(&a1, &b0);
let c1: Vec<i64> = itertools::zip_eq(c1_l, c1_r).map(|(l, r)| l + r).collect();
let c2: Vec<i64> = naive_mul(&a1, &b1);
// scale down, module Q, so result is \in R_q
let c0: Rq<Q, N> = arithmetic::ring::mul_div_round::<Q, N>(c0, T, Q);
let c1: Rq<Q, N> = arithmetic::ring::mul_div_round::<Q, N>(c1, T, Q);
let c2: Rq<Q, N> = arithmetic::ring::mul_div_round::<Q, N>(c2, T, Q);
(c0, c1, c2)
}
fn tensor_DBG<const PQ: u64, const T: u64>(
a: &Self,
b: &Self,
) -> (Rq<Q, N>, Rq<Q, N>, Rq<Q, N>) {
// iacr 2021-204:
// expand Q->PQ
// let a: RLWE<PQ, N> = a.remodule::<PQ>();
// let b: RLWE<PQ, N> = b.remodule::<PQ>();
// tensor
let c0: Rq<Q, N> = a.0 * b.0; // NTT mul
let c1: Rq<Q, N> = a.0 * b.1 + a.1 * b.0; // NTT mul
let c2: Rq<Q, N> = a.1 * b.1; // NTT mul
// scale down
let c0: Rq<Q, N> = c0.mul_div_round(T, Q);
let c1: Rq<Q, N> = c1.mul_div_round(T, Q);
let c2: Rq<Q, N> = c2.mul_div_round(T, Q);
// expand^-1 PQ->Q
// let c0: Rq<Q, N> = c0.remodule::<Q>();
// let c1: Rq<Q, N> = c1.remodule::<Q>();
// let c2: Rq<Q, N> = c2.remodule::<Q>();
(c0, c1, c2)
}
/// ciphertext multiplication
fn mul<const PQ: u64, const T: u64>(rlk: &RLK<PQ, N>, a: &Self, b: &Self) -> Self {
let (c0, c1, c2) = Self::tensor::<PQ, T>(a, b);
BFV::<Q, N, T>::relinearize_204::<PQ>(&rlk, &c0, &c1, &c2)
}
}
// naive mul in the ring Rq, reusing the ring::naive_mul and then applying mod(X^N +1)
fn tmp_naive_mul<const Q: u64, const N: usize>(a: Rq<Q, N>, b: Rq<Q, N>) -> Rq<Q, N> {
Rq::<Q, N>::from_vec_i64(arithmetic::ring::naive_mul(&a.to_r(), &b.to_r()))
} }
impl<const Q: u64, const N: usize> ops::Add<RLWE<Q, N>> for RLWE<Q, N> { impl<const Q: u64, const N: usize> ops::Add<RLWE<Q, N>> for RLWE<Q, N> {
@@ -39,17 +113,15 @@ impl<const Q: u64, const N: usize> ops::Add<RLWE<Q, N>> for RLWE<Q, N> {
} }
} }
impl<const Q: u64, const N: usize, const T: u64> ops::Add<&PR<T, N>> for &RLWE<Q, N> { impl<const Q: u64, const N: usize, const T: u64> ops::Add<&Rq<T, N>> for &RLWE<Q, N> {
type Output = RLWE<Q, N>; type Output = RLWE<Q, N>;
fn add(self, rhs: &PR<T, N>) -> Self::Output { fn add(self, rhs: &Rq<T, N>) -> Self::Output {
// todo!()
BFV::<Q, N, T>::add_const(self, rhs) BFV::<Q, N, T>::add_const(self, rhs)
} }
} }
impl<const Q: u64, const N: usize, const T: u64> ops::Mul<&PR<T, N>> for &RLWE<Q, N> { impl<const Q: u64, const N: usize, const T: u64> ops::Mul<&Rq<T, N>> for &RLWE<Q, N> {
type Output = RLWE<Q, N>; type Output = RLWE<Q, N>;
fn mul(self, rhs: &PR<T, N>) -> Self::Output { fn mul(self, rhs: &Rq<T, N>) -> Self::Output {
// todo!()
BFV::<Q, N, T>::mul_const(&self, rhs) BFV::<Q, N, T>::mul_const(&self, rhs)
} }
} }
@@ -71,50 +143,175 @@ impl<const Q: u64, const N: usize, const T: u64> BFV<Q, N, T> {
let Xi_err = Normal::new(0_f64, ERR_SIGMA)?; let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
// secret key // secret key
let s = PR::<Q, N>::rand_u64(&mut rng, Xi_key)?; let s = Rq::<Q, N>::rand_u64(&mut rng, Xi_key)?;
#[cfg(test)] // sanity check
assert!(s.infinity_norm() <= 1, "{:?}", s.coeffs());
// pk = (-a * s + e, a) // pk = (-a * s + e, a)
let a = PR::<Q, N>::rand_u64(&mut rng, Uniform::new(0_u64, Q))?; let a = Rq::<Q, N>::rand_u64(&mut rng, Uniform::new(0_u64, Q))?;
let e = PR::<Q, N>::rand_f64(&mut rng, Xi_err)?; let e = Rq::<Q, N>::rand_f64(&mut rng, Xi_err)?;
let pk: PublicKey<Q, N> = PublicKey((&(-a) * &s) + e, a.clone()); let pk: PublicKey<Q, N> = PublicKey((&(-a) * &s) + e, a.clone());
Ok((SecretKey(s), pk)) Ok((SecretKey(s), pk))
} }
pub fn encrypt(mut rng: impl Rng, pk: &PublicKey<Q, N>, m: &PR<T, N>) -> Result<RLWE<Q, N>> { pub fn encrypt(mut rng: impl Rng, pk: &PublicKey<Q, N>, m: &Rq<T, N>) -> Result<RLWE<Q, N>> {
let Xi_key = Uniform::new(-1_f64, 1_f64); // let Xi_key = Uniform::new(-1_f64, 1_f64);
// let Xi_key = Uniform::new(0_f64, 2_f64);
let Xi_key = Uniform::new(0_u64, 2_u64);
let Xi_err = Normal::new(0_f64, ERR_SIGMA)?; let Xi_err = Normal::new(0_f64, ERR_SIGMA)?;
let u = PR::<Q, N>::rand_f64(&mut rng, Xi_key)?; let u = Rq::<Q, N>::rand_u64(&mut rng, Xi_key)?;
let e_1 = PR::<Q, N>::rand_f64(&mut rng, Xi_err)?; let e_1 = Rq::<Q, N>::rand_f64_abs(&mut rng, Xi_err)?;
let e_2 = PR::<Q, N>::rand_f64(&mut rng, Xi_err)?; let e_2 = Rq::<Q, N>::rand_f64_abs(&mut rng, Xi_err)?;
// println!("{:?}", &e_1.coeffs());
// println!("{:?}", &e_2.coeffs());
#[cfg(test)] // sanity check
assert!(u.infinity_norm() <= 1, "{:?}", u.coeffs());
// migrate m's coeffs to the bigger modulus Q (from T) // migrate m's coeffs to the bigger modulus Q (from T)
let m = PR::<Q, N>::from_vec_u64(m.coeffs().iter().map(|m_i| m_i.0).collect()); let m = m.remodule::<Q>();
let c0 = &pk.0 * &u + e_1 + m * Self::DELTA; #[cfg(test)]
{
// sanity check // TODO rm
let m_remod_naive =
Rq::<Q, N>::from_vec_u64(m.coeffs().iter().map(|m_i| m_i.0).collect());
assert_eq!(m_remod_naive, m);
}
// let c0 = &pk.0 * &u + e_1 + m * Self::DELTA;
let c0 = &pk.0 * &u + e_1 + m.mul_div_round(Q, T); // TODO use DELTA?
let c1 = &pk.1 * &u + e_2; let c1 = &pk.1 * &u + e_2;
// let c0 = tmp_naive_mul(pk.0, u) + e_1 + m * Self::DELTA;
// let c0 = tmp_naive_mul(pk.0, u) + e_1 + m.mul_div_round(Q, T);
// let c1 = tmp_naive_mul(pk.1, u)
// // &pk.1 * &u
// + e_2;
Ok(RLWE::<Q, N>(c0, c1)) Ok(RLWE::<Q, N>(c0, c1))
} }
pub fn decrypt(sk: &SecretKey<Q, N>, c: &RLWE<Q, N>) -> PR<T, N> { pub fn decrypt(sk: &SecretKey<Q, N>, c: &RLWE<Q, N>) -> Rq<T, N> {
let cs = c.0 + c.1 * sk.0; // done in mod q let cs = c.0 + c.1 * sk.0; // done in mod q
let r: Vec<u64> = cs
.coeffs() // let c1s = tmp_naive_mul(c.1, sk.0);
.iter() // // let c1s = arithmetic::ring::naive_mul(&c.1.to_r(), &sk.0.to_r()); // TODO rm
.map(|e| ((T as f64 * e.0 as f64) / Q as f64).round() as u64) // // let c1s = Rq::<Q, N>::from_vec_i64(c1s);
.collect(); // let cs = c.0 + c1s;
PR::<T, N>::from_vec_u64(r)
// let r: Vec<u64> = cs
// .coeffs()
// .iter()
// .map(|e| ((T as f64 * e.0 as f64) / Q as f64).round() as u64)
// .collect();
// Rq::<T, N>::from_vec_u64(r)
let r: Rq<Q, N> = cs.mul_div_round(T, Q);
r.remodule::<T>()
} }
fn add_const(c: &RLWE<Q, N>, m: &PR<T, N>) -> RLWE<Q, N> { fn add_const(c: &RLWE<Q, N>, m: &Rq<T, N>) -> RLWE<Q, N> {
// assuming T<Q, move m from Zq<T> to Zq<Q> // assuming T<Q, move m from Zq<T> to Zq<Q>
let m = m.remodule::<Q>(); let m = m.remodule::<Q>();
RLWE::<Q, N>(c.0 + m * Self::DELTA, c.1) RLWE::<Q, N>(c.0 + m * Self::DELTA, c.1)
} }
fn mul_const(c: &RLWE<Q, N>, m: &PR<T, N>) -> RLWE<Q, N> { fn mul_const(c: &RLWE<Q, N>, m: &Rq<T, N>) -> RLWE<Q, N> {
// assuming T<Q, move m from Zq<T> to Zq<Q> // assuming T<Q, move m from Zq<T> to Zq<Q>
let m = m.remodule::<Q>(); let m = m.remodule::<Q>();
RLWE::<Q, N>(c.0 * m * Self::DELTA, c.1) RLWE::<Q, N>(c.0 * m * Self::DELTA, c.1)
} }
fn rlk_key<const PQ: u64>(mut rng: impl Rng, s: &SecretKey<Q, N>) -> Result<RLK<PQ, N>> {
let Xi_err = Normal::new(0_f64, ERR_SIGMA)?; // TODO review Xi' instead of Xi
let s = s.0.remodule::<PQ>();
let a = Rq::<PQ, N>::rand_u64(&mut rng, Uniform::new(0_u64, PQ))?;
let e = Rq::<PQ, N>::rand_f64(&mut rng, Xi_err)?;
// let rlk_1: Rq<PQ, N> = (&(-a) * &s) - e + (s * s) * P;
let P = PQ / Q;
// let rlk: RLK<PQ, N> = RLK::<PQ, N>((&(-a) * &s) - e + (s * s) * P, a.clone());
let rlk: RLK<PQ, N> = RLK::<PQ, N>(
-(tmp_naive_mul(a, s) + e) + tmp_naive_mul(s, s) * P,
a.clone(),
);
// let rlk: RLK<PQ, N> = RLK::<PQ, N>(-(&a * &s + e) + (s * s) * P, a.clone());
// let pq = P * Q;
// let a = Rq::<Q, N>::rand_u64(&mut rng, Uniform::new(0_u64, pq))?;
// let e = Rq::<Q, N>::rand_f64(&mut rng, Xi_err)?;
//
// let rlk_0: Rq<Q, N> = (&(-a) * &s) - e + (s * s) * P;
// let rlk_0 = rlk_0.remodule::<>();
// let rlk: RLK<Q, N> = RLK(rlk_0, a);
Ok(rlk)
}
fn relinearize<const PQ: u64>(
rlk: &RLK<PQ, N>,
c0: &Rq<Q, N>,
c1: &Rq<Q, N>,
c2: &Rq<Q, N>,
) -> RLWE<Q, N> {
let P = PQ / Q;
// let c2 = c2.remodule::<PQ>();
// let c2 = c2.to_r();
let c2rlk0: Vec<f64> = (c2.to_r() * rlk.0.to_r())
.coeffs()
.iter()
.map(|e| (*e as f64 / P as f64).round())
.collect();
let c2rlk1: Vec<f64> = (c2.to_r() * rlk.1.to_r())
.coeffs()
.iter()
.map(|e| (*e as f64 / P as f64).round())
.collect();
let r0 = Rq::<Q, N>::from_vec_f64(c2rlk0);
let r1 = Rq::<Q, N>::from_vec_f64(c2rlk1);
let res = RLWE::<Q, N>(c0 + &r0, c1 + &r1);
res
}
fn relinearize_204<const PQ: u64>(
rlk: &RLK<PQ, N>,
c0: &Rq<Q, N>,
c1: &Rq<Q, N>,
c2: &Rq<Q, N>,
) -> RLWE<Q, N> {
let P = PQ / Q;
// let c2 = c2.remodule::<PQ>();
// let c2 = c2.to_r();
// let c2rlk0: Vec<f64> = (c2.remodule::<PQ>() * rlk.0)
use arithmetic::ring::naive_mul;
let c2rlk0: Vec<i64> = naive_mul(&c2.to_r(), &rlk.0.to_r());
// .coeffs()
// .iter()
// .map(|e| (*e as f64 / P as f64).round())
// .collect();
let c2rlk1: Vec<i64> = naive_mul(&c2.to_r(), &rlk.1.to_r());
// .coeffs()
// .iter()
// .map(|e| (*e as f64 / P as f64).round())
// .collect();
//
let r0: Rq<Q, N> = arithmetic::ring::mul_div_round::<Q, N>(c2rlk0, 1, P);
let r1: Rq<Q, N> = arithmetic::ring::mul_div_round::<Q, N>(c2rlk1, 1, P);
// let r0 = Rq::<Q, N>::from_vec_f64(c2rlk0);
// let r1 = Rq::<Q, N>::from_vec_f64(c2rlk1);
let res = RLWE::<Q, N>(c0 + &r0, c1 + &r1);
res
}
} }
#[cfg(test)] #[cfg(test)]
@@ -133,15 +330,17 @@ mod tests {
let mut rng = rand::thread_rng(); let mut rng = rand::thread_rng();
for _ in 0..1000 {
let (sk, pk) = S::new_key(&mut rng)?; let (sk, pk) = S::new_key(&mut rng)?;
let msg_dist = Uniform::new(0_u64, T); let msg_dist = Uniform::new(0_u64, T);
let m = PR::<T, N>::rand_u64(&mut rng, msg_dist)?; let m = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let c = S::encrypt(rng, &pk, &m)?; let c = S::encrypt(&mut rng, &pk, &m)?;
let m_recovered = S::decrypt(&sk, &c); let m_recovered = S::decrypt(&sk, &c);
assert_eq!(m, m_recovered); assert_eq!(m, m_recovered);
}
Ok(()) Ok(())
} }
@@ -149,7 +348,8 @@ mod tests {
#[test] #[test]
fn test_addition() -> Result<()> { fn test_addition() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1; const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 32; // const N: usize = 32;
const N: usize = 4;
const T: u64 = 4; // plaintext modulus const T: u64 = 4; // plaintext modulus
type S = BFV<Q, N, T>; type S = BFV<Q, N, T>;
@@ -158,8 +358,8 @@ mod tests {
let (sk, pk) = S::new_key(&mut rng)?; let (sk, pk) = S::new_key(&mut rng)?;
let msg_dist = Uniform::new(0_u64, T); let msg_dist = Uniform::new(0_u64, T);
let m1 = PR::<T, N>::rand_u64(&mut rng, msg_dist)?; let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2 = PR::<T, N>::rand_u64(&mut rng, msg_dist)?; let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let c1 = S::encrypt(&mut rng, &pk, &m1)?; let c1 = S::encrypt(&mut rng, &pk, &m1)?;
let c2 = S::encrypt(&mut rng, &pk, &m2)?; let c2 = S::encrypt(&mut rng, &pk, &m2)?;
@@ -174,7 +374,7 @@ mod tests {
} }
#[test] #[test]
fn test_constant_add_mul() -> Result<()> { fn test_constant_add() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1; const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 32; const N: usize = 32;
const T: u64 = 4; // plaintext modulus const T: u64 = 4; // plaintext modulus
@@ -185,66 +385,271 @@ mod tests {
let (sk, pk) = S::new_key(&mut rng)?; let (sk, pk) = S::new_key(&mut rng)?;
let msg_dist = Uniform::new(0_u64, T); let msg_dist = Uniform::new(0_u64, T);
let m1 = PR::<T, N>::rand_u64(&mut rng, msg_dist)?; let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2_const = PR::<T, N>::rand_u64(&mut rng, msg_dist)?; let m2_const = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let c1 = S::encrypt(&mut rng, &pk, &m1)?; let c1 = S::encrypt(&mut rng, &pk, &m1)?;
let c3_add = &c1 + &m2_const; let c3_add = &c1 + &m2_const;
let c3_mul = &c1 * &m2_const; // let c3_mul = &c1 * &m2_const;
let m3_add_recovered = S::decrypt(&sk, &c3_add); let m3_add_recovered = S::decrypt(&sk, &c3_add);
let m3_mul_recovered = S::decrypt(&sk, &c3_mul); // let m3_mul_recovered = S::decrypt(&sk, &c3_mul);
assert_eq!(m1 + m2_const, m3_add_recovered); assert_eq!(m1 + m2_const, m3_add_recovered);
//
let mut mul_res = naive_poly_mul::<T>(&m1.coeffs().to_vec(), &m2_const.coeffs().to_vec()); // let mut mul_res = naive_poly_mul::<T>(&m1.coeffs().to_vec(), &m2_const.coeffs().to_vec());
arithmetic::ring::modulus::<T, N>(&mut mul_res); // arithmetic::ring::modulus::<T, N>(&mut mul_res);
dbg!(&mul_res); // dbg!(&mul_res);
let mul_res_2 = // let mul_res_2 =
naive_poly_mul_2::<T, N>(&m1.coeffs().to_vec(), &m2_const.coeffs().to_vec()); // naive_poly_mul_2::<T, N>(&m1.coeffs().to_vec(), &m2_const.coeffs().to_vec());
assert_eq!(mul_res, mul_res_2); // assert_eq!(mul_res, mul_res_2);
let mul_res = PR::<T, N>::from_vec(mul_res); // let mul_res = PR::<T, N>::from_vec(mul_res);
assert_eq!(mul_res.coeffs(), m3_mul_recovered.coeffs()); // assert_eq!(mul_res.coeffs(), m3_mul_recovered.coeffs());
Ok(()) Ok(())
} }
fn naive_poly_mul<const T: u64>(a: &[Zq<T>], b: &[Zq<T>]) -> Vec<Zq<T>> { #[test]
let mut result: Vec<Zq<T>> = vec![Zq::zero(); a.len() + b.len() - 1]; fn test_tensor() -> Result<()> {
for (i, &ai) in a.iter().enumerate() { const Q: u64 = 2u64.pow(16) + 1; // q prime, and 2^q + 1 shape
for (j, &bj) in b.iter().enumerate() { const N: usize = 8;
result[i + j] = result[i + j] + (ai * bj); const T: u64 = 4; // plaintext modulus
}
}
result
}
fn naive_poly_mul_2<const T: u64, const N: usize>(
poly1: &[Zq<T>],
poly2: &[Zq<T>],
) -> Vec<Zq<T>> {
let degree1 = poly1.len();
let degree2 = poly2.len();
// The degree of the resulting polynomial will be degree1 + degree2 - 1 // const P: u64 = Q;
let mut result = vec![Zq::zero(); degree1 + degree2 - 1]; const P: u64 = Q * Q;
// const P: u64 = 2_u64.pow(13) * Q + 1;
const PQ: u64 = P * Q;
// Perform the multiplication let mut rng = rand::thread_rng();
for i in 0..degree1 {
for j in 0..degree2 { let msg_dist = Uniform::new(0_u64, T);
result[i + j] = result[i + j] + poly1[i] * poly2[j]; for _ in 0..10_000 {
} let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
test_tensor_opt::<Q, N, T, PQ>(&mut rng, m1, m2)?;
} }
// Reduce the result modulo x^N + 1 Ok(())
let mut reduced_result = vec![Zq::zero(); N]; }
fn test_tensor_opt<const Q: u64, const N: usize, const T: u64, const PQ: u64>(
mut rng: impl Rng,
m1: Rq<T, N>,
m2: Rq<T, N>,
) -> Result<()> {
let (sk, pk) = BFV::<Q, N, T>::new_key(&mut rng)?;
for i in 0..result.len() { let c1 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m1)?;
let mod_index = i % N; // wrap around using modulo N let c2 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m2)?;
reduced_result[mod_index] += result[i];
let (c_a, c_b, c_c) = RLWE::<Q, N>::tensor::<PQ, T>(&c1, &c2);
// decrypt non-relinearized mul result
let m3: Rq<Q, N> = c_a + c_b * sk.0 + c_c * sk.0 * sk.0;
let m3: Rq<Q, N> = m3.mul_div_round(T, Q);
let m3 = m3.remodule::<T>();
let naive = (m1.to_r() * m2.to_r()).to_rq::<T>();
assert_eq!(
m3.coeffs().to_vec(),
naive.coeffs().to_vec(),
"\n\nfor testing:\nlet m1 = Rq::<T, N>::from_vec_u64(vec!{:?});\nlet m2 = Rq::<T, N>::from_vec_u64(vec!{:?});\n",
m1.coeffs(),
m2.coeffs()
);
if m3.coeffs().to_vec() != naive.coeffs().to_vec() {
return Err(anyhow!("not eq"));
} }
// Return the reduced polynomial Ok(())
reduced_result }
fn test_tensor_opt_DBG<const Q: u64, const N: usize, const T: u64, const PQ: u64>(
mut rng: impl Rng,
m1: Rq<T, N>,
m2: Rq<T, N>,
) -> Result<()> {
let (sk, pk) = BFV::<Q, N, T>::new_key(&mut rng)?;
let c1 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m1)?;
let c2 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m2)?;
let (c_a, c_b, c_c) = RLWE::<Q, N>::tensor::<PQ, T>(&c1, &c2);
// decrypt non-relinearized mul result
let m3: Rq<Q, N> = c_a + c_b * sk.0 + c_c * sk.0 * sk.0;
dbg!(m3);
let m3: Rq<Q, N> = m3.mul_div_round(T, Q);
dbg!(m3);
let m3 = m3.remodule::<T>();
dbg!(m3);
// let naive = (m1.to_r() * m2.to_r()).to_rq::<T>();
// let naive = m1.remodule::<Q>() * m2.remodule::<Q>();
let naive = (m1.remodule::<Q>() * m2.remodule::<Q>()).remodule::<T>();
dbg!(naive);
assert_eq!(
m3.coeffs().to_vec(),
naive.coeffs().to_vec(),
"\n\nfor testing:\nlet m1 = Rq::<T, N>::from_vec_u64(vec!{:?});\nlet m2 = Rq::<T, N>::from_vec_u64(vec!{:?});\n",
m1.coeffs(),
m2.coeffs()
);
// if m3.coeffs().to_vec() != naive.coeffs().to_vec() {
// return Err(anyhow!("not eq"));
// }
Ok(())
}
#[test]
fn test_mul_relin() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 32;
const T: u64 = 4; // plaintext modulus
type S = BFV<Q, N, T>;
const P: u64 = Q * Q;
const PQ: u64 = P * Q;
let mut rng = rand::thread_rng();
let msg_dist = Uniform::new(0_u64, T);
for _ in 0..100 {
let m1 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let m2 = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
test_mul_relin_opt::<Q, N, T, PQ>(&mut rng, m1, m2)?;
}
Ok(())
}
fn test_mul_relin_opt<const Q: u64, const N: usize, const T: u64, const PQ: u64>(
mut rng: impl Rng,
m1: Rq<T, N>,
m2: Rq<T, N>,
) -> Result<()> {
let (sk, pk) = BFV::<Q, N, T>::new_key(&mut rng)?;
let rlk = BFV::<Q, N, T>::rlk_key::<PQ>(&mut rng, &sk)?;
let c1 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m1)?;
let c2 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &m2)?;
let c3 = RLWE::<Q, N>::mul::<PQ, T>(&rlk, &c1, &c2);
let m3 = BFV::<Q, N, T>::decrypt(&sk, &c3);
let naive = (m1.to_r() * m2.to_r()).to_rq::<T>();
assert_eq!(m3.coeffs().to_vec(), naive.coeffs().to_vec(),
"\n\nfor testing:\nlet m1 = Rq::<T, N>::from_vec_u64(vec!{:?});\nlet m2 = Rq::<T, N>::from_vec_u64(vec!{:?});\n",
m1.coeffs(),
m2.coeffs()
);
Ok(())
}
#[test]
fn test_naive_mul() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1; // prime, and 2^q + 1 shape
const N: usize = 4;
const T: u64 = 4; // plaintext modulus
let mut rng = rand::thread_rng();
let msg_dist = Uniform::new(0_u64, T);
// for _ in 0..10_000 {
for _ in 0..2 {
println!("---");
// let a = Rq::<Q, N>::rand_u64(&mut rng, msg_dist)?;
// let b = Rq::<Q, N>::rand_u64(&mut rng, msg_dist)?;
// let a = Rq::<Q, N>::from_vec_u64(vec![Q - 1, Q - 2, Q - 3, Q - 3]);
// let b = Rq::<Q, N>::from_vec_u64(vec![Q - 3, Q - 3, Q - 2, Q - 1]);
let a = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
let b = Rq::<T, N>::rand_u64(&mut rng, msg_dist)?;
dbg!(&a);
dbg!(&b);
let (_, pk) = BFV::<Q, N, T>::new_key(&mut rng)?;
let ciph1 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &a)?;
let ciph2 = BFV::<Q, N, T>::encrypt(&mut rng, &pk, &b)?;
let a = ciph1.0;
let b = ciph2.0;
dbg!(&a);
dbg!(&b);
let c0: Vec<i64> = arithmetic::ring::naive_mul(&a.to_r(), &b.to_r());
let c0 = Rq::<Q, N>::from_vec_i64(c0);
let c1 = tmp_naive_mul(a, b); // naive mul
let c2: Rq<Q, N> = a * b; // NTT mul
println!("{:?}", c0.coeffs());
println!("{:?}", c1.coeffs());
println!("{:?}", c2.coeffs());
assert_eq!(c0, c2);
assert_eq!(c1, c2);
// scale by Delta=Q/T
let a = a.mul_div_round(Q, T);
let b = b.mul_div_round(Q, T);
dbg!(&a);
dbg!(&b);
let c0: Vec<i64> = arithmetic::ring::naive_mul(&a.to_r(), &b.to_r());
let c0 = Rq::<Q, N>::from_vec_i64(c0);
let c1 = tmp_naive_mul(a, b); // naive mul
let c2: Rq<Q, N> = a * b; // NTT mul
println!("{:?}", c0.coeffs());
println!("{:?}", c1.coeffs());
println!("{:?}", c2.coeffs());
assert_eq!(c0, c2);
assert_eq!(c1, c2);
let c0 = c0.mul_div_round(T, Q);
let c1 = c1.mul_div_round(T, Q);
let c2 = c2.mul_div_round(T, Q);
println!("{:?}", c0.coeffs());
println!("{:?}", c1.coeffs());
println!("{:?}", c2.coeffs());
assert_eq!(c0, c2);
assert_eq!(c1, c2);
/*
// now same as before, but multiplying by T/Q
let c0: Vec<i64> = arithmetic::ring::naive_mul(&a.to_r(), &b.to_r());
let c0: Vec<f64> = c0
.iter()
.map(|e| ((T as f64 * *e as f64) / Q as f64).round())
.collect();
let c0 = Rq::<Q, N>::from_vec_f64(c0);
dbg!(&c0.coeffs());
let a = a.mul_div_round(T, Q);
let b = b.mul_div_round(T, Q);
println!("a{:?}", a.coeffs());
println!("b{:?}", b.coeffs());
let c4: Vec<i64> = arithmetic::ring::naive_mul(&a.to_r(), &b.to_r());
let c4 = Rq::<Q, N>::from_vec_i64(c4);
let c4 = c4.mul_div_round(T, Q);
let c1 = tmp_naive_mul(a, b); // naive mul
let c1 = c1.mul_div_round(T, Q);
let c2 = a * b; // NTT mul
let c2 = c2.mul_div_round(T, Q);
println!("{:?}", c0.coeffs());
println!("{:?}", c1.coeffs());
println!("{:?}", c2.coeffs());
println!("{:?}", c4.coeffs());
assert_eq!(c0, c2);
assert_eq!(c4, c2);
assert_eq!(c1, c2);
*/
}
Ok(())
} }
} }