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

Rm const generics (#2)

Browse Source 
* arith: get rid of constant generics. Reason:

using constant generics was great for allocating the arrays in the
stack, which is faster, but when started to use bigger parameter values,
in some cases it was overflowing the stack. This commit removes all the
constant generics in all of the `arith` crate, which in some cases slows
a bit the performance, but allows for bigger parameter values (on the
ones that affect lengths, like N and K).

* bfv: get rid of constant generics (reason in previous commit)

* ckks: get rid of constant generics (reason in two commits ago)

* group ring params under a single struct

* gfhe: get rid of constant generics

* tfhe: get rid of constant generics

* polish & clean a bit

* add methods for encoding constants for ct-pt-multiplication
This commit is contained in:
arnaucube
2025-08-14 18:32:43 +02:00
committed by GitHub
parent 13abadf6e1
commit fb1fb6b4e9
23 changed files with 2737 additions and 1887 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
arith/Cargo.toml
View File

@@ -8,6 +8,7 @@ anyhow = { workspace = true }
rand = { workspace = true }
rand_distr = { workspace = true }
itertools = { workspace = true }
lazy_static = "1.5.0"
# 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

5
arith/src/lib.rs
View File

@@ -15,17 +15,16 @@ pub mod ring_nq;
pub mod ring_torus;
pub mod tuple_ring;
mod naive_ntt; // note: for dev only
// mod naive_ntt; // note: for dev only
pub mod ntt;
// expose objects
pub use complex::C;
pub use matrix::Matrix;
pub use torus::T64;
pub use zq::Zq;
pub use ring::Ring;
pub use ring::{Ring, RingParam};
pub use ring_n::R;
pub use ring_nq::Rq;
pub use ring_torus::Tn;

178
arith/src/ntt.rs
View File

@@ -1,34 +1,60 @@
//! Implementation of the NTT & iNTT, following the CT & GS algorighms, more details in
//! https://eprint.iacr.org/2017/727.pdf, some notes at
//! https://github.com/arnaucube/math/blob/master/notes_ntt.pdf .
use crate::zq::Zq;
//!
//! NOTE: initially I implemented it with fixed Q & N, given as constant
//! generics; but once using real-world parameters, the stack could not handle
//! it, so moved to use Vec instead of fixed-sized arrays, and adapted the NTT
//! implementation to that too.
use crate::{ring::RingParam, ring_nq::Rq, zq::Zq};
use std::collections::HashMap;
#[derive(Debug)]
pub struct NTT<const Q: u64, const N: usize> {}
pub struct NTT {}
impl<const Q: u64, const N: usize> NTT<Q, N> {
const N_INV: Zq<Q> = Zq(const_inv_mod::<Q>(N as u64));
// since we work over Zq[X]/(X^N+1) (negacyclic), get the 2*N-th root of unity
pub(crate) const ROOT_OF_UNITY: u64 = primitive_root_of_unity::<Q>(2 * N);
pub(crate) const ROOTS_OF_UNITY: [Zq<Q>; N] = roots_of_unity(Self::ROOT_OF_UNITY);
const ROOTS_OF_UNITY_INV: [Zq<Q>; N] = roots_of_unity_inv(Self::ROOTS_OF_UNITY);
use std::sync::{Mutex, OnceLock};
static CACHE: OnceLock<Mutex<HashMap<(u64, usize), (Vec<Zq>, Vec<Zq>, Zq)>>> = OnceLock::new();
fn roots(q: u64, n: usize) -> (Vec<Zq>, Vec<Zq>, Zq) {
let cache_lock = CACHE.get_or_init(|| Mutex::new(HashMap::new()));
let mut cache = cache_lock.lock().unwrap();
if let Some(value) = cache.get(&(q, n)) {
return value.clone();
}
let n_inv: Zq = Zq {
q,
v: const_inv_mod(q, n as u64),
};
let root_of_unity: u64 = primitive_root_of_unity(q, 2 * n);
let roots_of_unity: Vec<Zq> = roots_of_unity(q, n, root_of_unity);
let roots_of_unity_inv: Vec<Zq> = roots_of_unity_inv(q, n, roots_of_unity.clone());
let value = (roots_of_unity, roots_of_unity_inv, n_inv);
cache.insert((q, n), value.clone());
value
}
impl<const Q: u64, const N: usize> NTT<Q, N> {
impl NTT {
/// implements the Cooley-Tukey (CT) algorithm. Details at
/// https://eprint.iacr.org/2017/727.pdf, also some notes at section 3.1 of
/// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
pub fn ntt(a: [Zq<Q>; N]) -> [Zq<Q>; N] {
let mut t = N / 2;
pub fn ntt(a: &Rq) -> Rq {
let (q, n) = (a.param.q, a.param.n);
let (roots_of_unity, _, _) = roots(q, n);
let mut t = n / 2;
let mut m = 1;
let mut r: [Zq<Q>; N] = a.clone();
while m < N {
let mut r: Vec<Zq> = a.coeffs.clone();
while m < n {
let mut k = 0;
for i in 0..m {
let S: Zq<Q> = Self::ROOTS_OF_UNITY[m + i];
let S: Zq = roots_of_unity[m + i];
for j in k..k + t {
let U: Zq<Q> = r[j];
let V: Zq<Q> = r[j + t] * S;
let U: Zq = r[j];
let V: Zq = r[j + t] * S;
r[j] = U + V;
r[j + t] = U - V;
}
@@ -37,23 +63,32 @@ impl<const Q: u64, const N: usize> NTT<Q, N> {
t /= 2;
m *= 2;
}
r
// TODO think if maybe not return a Rq type, or if returned Rq, maybe
// fill the `evals` field, which is what we're actually returning here
Rq {
param: RingParam { q, n },
coeffs: r,
evals: None,
}
}
/// implements the Cooley-Tukey (CT) algorithm. Details at
/// https://eprint.iacr.org/2017/727.pdf, also some notes at section 3.2 of
/// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
pub fn intt(a: [Zq<Q>; N]) -> [Zq<Q>; N] {
pub fn intt(a: &Rq) -> Rq {
let (q, n) = (a.param.q, a.param.n);
let (_, roots_of_unity_inv, n_inv) = roots(q, n);
let mut t = 1;
let mut m = N / 2;
let mut r: [Zq<Q>; N] = a.clone();
let mut m = n / 2;
let mut r: Vec<Zq> = a.coeffs.clone();
while m > 0 {
let mut k = 0;
for i in 0..m {
let S: Zq<Q> = Self::ROOTS_OF_UNITY_INV[m + i];
let S: Zq = roots_of_unity_inv[m + i];
for j in k..k + t {
let U: Zq<Q> = r[j];
let V: Zq<Q> = r[j + t];
let U: Zq = r[j];
let V: Zq = r[j + t];
r[j] = U + V;
r[j + t] = (U - V) * S;
}
@@ -62,26 +97,32 @@ impl<const Q: u64, const N: usize> NTT<Q, N> {
t *= 2;
m /= 2;
}
for i in 0..N {
r[i] = r[i] * Self::N_INV;
for i in 0..n {
r[i] = r[i] * n_inv;
}
Rq {
param: RingParam { q, n },
coeffs: r,
// TODO maybe at `evals` place the inputed `a` which is the evals
// format
evals: None,
}
r
}
}
/// computes a primitive N-th root of unity using the method described by Thomas
/// Pornin in https://crypto.stackexchange.com/a/63616
const fn primitive_root_of_unity<const Q: u64>(N: usize) -> u64 {
assert!(N.is_power_of_two());
assert!((Q - 1) % N as u64 == 0);
const fn primitive_root_of_unity(q: u64, n: usize) -> u64 {
assert!(n.is_power_of_two());
assert!((q - 1) % n as u64 == 0);
let n_u64 = n as u64;
let n: u64 = N as u64;
let mut k = 1;
while k < Q {
while k < q {
// alternatively could get a random k at each iteration, if so, add the following if:
// `if k == 0 { continue; }`
let w = const_exp_mod::<Q>(k, (Q - 1) / n);
if const_exp_mod::<Q>(w, n / 2) != 1 {
let w = const_exp_mod(q, k, (q - 1) / n_u64);
if const_exp_mod(q, w, n_u64 / 2) != 1 {
return w; // w is a primitive N-th root of unity
}
k += 1;
@@ -89,78 +130,85 @@ const fn primitive_root_of_unity<const Q: u64>(N: usize) -> u64 {
panic!("No primitive root of unity");
}
const fn roots_of_unity<const Q: u64, const N: usize>(w: u64) -> [Zq<Q>; N] {
let mut r: [Zq<Q>; N] = [Zq(0u64); N];
fn roots_of_unity(q: u64, n: usize, w: u64) -> Vec<Zq> {
let mut r: Vec<Zq> = vec![Zq { q, v: 0 }; n];
let mut i = 0;
let log_n = N.ilog2();
while i < N {
let log_n = n.ilog2();
while i < n {
// (return the roots in bit-reverset order)
let j = ((i as u64).reverse_bits() >> (64 - log_n)) as usize;
r[i] = Zq(const_exp_mod::<Q>(w, j as u64));
r[i] = Zq {
q,
v: const_exp_mod(q, w, j as u64),
};
i += 1;
}
r
}
const fn roots_of_unity_inv<const Q: u64, const N: usize>(v: [Zq<Q>; N]) -> [Zq<Q>; N] {
fn roots_of_unity_inv(q: u64, n: usize, v: Vec<Zq>) -> Vec<Zq> {
// assumes that the inputted roots are already in bit-reverset order
let mut r: [Zq<Q>; N] = [Zq(0u64); N];
let mut r: Vec<Zq> = vec![Zq { q, v: 0 }; n];
let mut i = 0;
while i < N {
r[i] = Zq(const_inv_mod::<Q>(v[i].0));
while i < n {
r[i] = Zq {
q,
v: const_inv_mod(q, v[i].v),
};
i += 1;
}
r
}
/// returns x^k mod Q
const fn const_exp_mod<const Q: u64>(x: u64, k: u64) -> u64 {
const fn const_exp_mod(q: u64, x: u64, k: u64) -> u64 {
// 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;
x = x % q as u128;
// exponentiation by square strategy
while k > 0 {
if k % 2 == 1 {
r = (r * x) % Q as u128;
r = (r * x) % q as u128;
}
x = (x * x) % Q as u128;
x = (x * x) % q as u128;
k /= 2;
}
r as u64
}
/// returns x^-1 mod Q
const fn const_inv_mod<const Q: u64>(x: u64) -> u64 {
const fn const_inv_mod(q: u64, x: u64) -> u64 {
// by Fermat's Little Theorem, x^-1 mod q \equiv x^{q-2} mod q
const_exp_mod::<Q>(x, Q - 2)
const_exp_mod(q, x, q - 2)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::Ring;
use anyhow::Result;
use std::array;
#[test]
fn test_ntt() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 4;
let q: u64 = 2u64.pow(16) + 1;
let n: usize = 4;
let param = RingParam { q, n };
let a: [u64; N] = [1u64, 2, 3, 4];
let a: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(a[i]));
let a: Vec<u64> = vec![1u64, 2, 3, 4];
let a: Rq = Rq::from_vec_u64(&param, a);
let a_ntt = NTT::<Q, N>::ntt(a);
let a_ntt = NTT::ntt(&a);
let a_intt = NTT::<Q, N>::intt(a_ntt);
let a_intt = NTT::intt(&a_ntt);
dbg!(&a);
dbg!(&a_ntt);
dbg!(&a_intt);
dbg!(NTT::<Q, N>::ROOT_OF_UNITY);
dbg!(NTT::<Q, N>::ROOTS_OF_UNITY);
// dbg!(NTT::ROOT_OF_UNITY);
// dbg!(NTT::ROOTS_OF_UNITY);
assert_eq!(a, a_intt);
Ok(())
@@ -168,18 +216,18 @@ mod tests {
#[test]
fn test_ntt_loop() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 512;
let q: u64 = 2u64.pow(16) + 1;
let n: usize = 512;
let param = RingParam { q, n };
use rand::distributions::Distribution;
use rand::distributions::Uniform;
let mut rng = rand::thread_rng();
let dist = Uniform::new(0_f64, Q as f64);
let dist = Uniform::new(0_f64, q as f64);
for _ in 0..100 {
let a: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng)));
let a_ntt = NTT::<Q, N>::ntt(a);
let a_intt = NTT::<Q, N>::intt(a_ntt);
for _ in 0..1000 {
let a: Rq = Rq::rand(&mut rng, dist, &param);
let a_ntt = NTT::ntt(&a);
let a_intt = NTT::intt(&a_ntt);
assert_eq!(a, a_intt);
}
Ok(())

187
arith/src/ntt_fixedsize.rs Normal file
View File

@@ -0,0 +1,187 @@
//! Implementation of the NTT & iNTT, following the CT & GS algorighms, more details in
//! https://eprint.iacr.org/2017/727.pdf, some notes at
//! https://github.com/arnaucube/math/blob/master/notes_ntt.pdf .
use crate::zq::Zq;
#[derive(Debug)]
pub struct NTT<const Q: u64, const N: usize> {}
impl<const Q: u64, const N: usize> NTT<Q, N> {
const N_INV: Zq<Q> = Zq(const_inv_mod::<Q>(N as u64));
// since we work over Zq[X]/(X^N+1) (negacyclic), get the 2*N-th root of unity
pub(crate) const ROOT_OF_UNITY: u64 = primitive_root_of_unity::<Q>(2 * N);
pub(crate) const ROOTS_OF_UNITY: [Zq<Q>; N] = roots_of_unity(Self::ROOT_OF_UNITY);
const ROOTS_OF_UNITY_INV: [Zq<Q>; N] = roots_of_unity_inv(Self::ROOTS_OF_UNITY);
}
impl<const Q: u64, const N: usize> NTT<Q, N> {
/// implements the Cooley-Tukey (CT) algorithm. Details at
/// https://eprint.iacr.org/2017/727.pdf, also some notes at section 3.1 of
/// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
pub fn ntt(a: [Zq<Q>; N]) -> [Zq<Q>; N] {
let mut t = N / 2;
let mut m = 1;
let mut r: [Zq<Q>; N] = a.clone();
while m < N {
let mut k = 0;
for i in 0..m {
let S: Zq<Q> = Self::ROOTS_OF_UNITY[m + i];
for j in k..k + t {
let U: Zq<Q> = r[j];
let V: Zq<Q> = r[j + t] * S;
r[j] = U + V;
r[j + t] = U - V;
}
k = k + 2 * t;
}
t /= 2;
m *= 2;
}
r
}
/// implements the Cooley-Tukey (CT) algorithm. Details at
/// https://eprint.iacr.org/2017/727.pdf, also some notes at section 3.2 of
/// https://github.com/arnaucube/math/blob/master/notes_ntt.pdf
pub fn intt(a: [Zq<Q>; N]) -> [Zq<Q>; N] {
let mut t = 1;
let mut m = N / 2;
let mut r: [Zq<Q>; N] = a.clone();
while m > 0 {
let mut k = 0;
for i in 0..m {
let S: Zq<Q> = Self::ROOTS_OF_UNITY_INV[m + i];
for j in k..k + t {
let U: Zq<Q> = r[j];
let V: Zq<Q> = r[j + t];
r[j] = U + V;
r[j + t] = (U - V) * S;
}
k += 2 * t;
}
t *= 2;
m /= 2;
}
for i in 0..N {
r[i] = r[i] * Self::N_INV;
}
r
}
}
/// computes a primitive N-th root of unity using the method described by Thomas
/// Pornin in https://crypto.stackexchange.com/a/63616
const fn primitive_root_of_unity<const Q: u64>(N: usize) -> u64 {
assert!(N.is_power_of_two());
assert!((Q - 1) % N as u64 == 0);
let n: u64 = N as u64;
let mut k = 1;
while k < Q {
// alternatively could get a random k at each iteration, if so, add the following if:
// `if k == 0 { continue; }`
let w = const_exp_mod::<Q>(k, (Q - 1) / n);
if const_exp_mod::<Q>(w, n / 2) != 1 {
return w; // w is a primitive N-th root of unity
}
k += 1;
}
panic!("No primitive root of unity");
}
const fn roots_of_unity<const Q: u64, const N: usize>(w: u64) -> [Zq<Q>; N] {
let mut r: [Zq<Q>; N] = [Zq(0u64); N];
let mut i = 0;
let log_n = N.ilog2();
while i < N {
// (return the roots in bit-reverset order)
let j = ((i as u64).reverse_bits() >> (64 - log_n)) as usize;
r[i] = Zq(const_exp_mod::<Q>(w, j as u64));
i += 1;
}
r
}
const fn roots_of_unity_inv<const Q: u64, const N: usize>(v: [Zq<Q>; N]) -> [Zq<Q>; N] {
// assumes that the inputted roots are already in bit-reverset order
let mut r: [Zq<Q>; N] = [Zq(0u64); N];
let mut i = 0;
while i < N {
r[i] = Zq(const_inv_mod::<Q>(v[i].0));
i += 1;
}
r
}
/// returns x^k mod Q
const fn const_exp_mod<const Q: u64>(x: u64, k: u64) -> u64 {
// 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 as u128;
}
x = (x * x) % Q as u128;
k /= 2;
}
r as u64
}
/// returns x^-1 mod Q
const fn const_inv_mod<const Q: u64>(x: u64) -> u64 {
// by Fermat's Little Theorem, x^-1 mod q \equiv x^{q-2} mod q
const_exp_mod::<Q>(x, Q - 2)
}
#[cfg(test)]
mod tests {
use super::*;
use anyhow::Result;
use std::array;
#[test]
fn test_ntt() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 4;
let a: [u64; N] = [1u64, 2, 3, 4];
let a: [Zq<Q>; N] = array::from_fn(|i| Zq::from_u64(a[i]));
let a_ntt = NTT::<Q, N>::ntt(a);
let a_intt = NTT::<Q, N>::intt(a_ntt);
dbg!(&a);
dbg!(&a_ntt);
dbg!(&a_intt);
dbg!(NTT::<Q, N>::ROOT_OF_UNITY);
dbg!(NTT::<Q, N>::ROOTS_OF_UNITY);
assert_eq!(a, a_intt);
Ok(())
}
#[test]
fn test_ntt_loop() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 512;
use rand::distributions::Distribution;
use rand::distributions::Uniform;
let mut rng = rand::thread_rng();
let dist = Uniform::new(0_f64, Q as f64);
for _ in 0..100 {
let a: [Zq<Q>; N] = array::from_fn(|_| Zq::from_f64(dist.sample(&mut rng)));
let a_ntt = NTT::<Q, N>::ntt(a);
let a_intt = NTT::<Q, N>::intt(a_ntt);
assert_eq!(a, a_intt);
}
Ok(())
}
}

22
arith/src/ring.rs
View File

@@ -3,6 +3,12 @@ use std::fmt::Debug;
use std::iter::Sum;
use std::ops::{Add, AddAssign, Mul, Neg, Sub, SubAssign};
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RingParam {
pub q: u64, // TODO think if really needed or it's fine with coeffs[0].q
pub n: usize,
}
/// Represents a ring element. Currently implemented by ring_nq.rs#Rq and
/// ring_torus.rs#Tn. Is not a 'pure algebraic ring', but more a custom trait
/// definition which includes methods like `mod_switch`.
@@ -21,27 +27,25 @@ pub trait Ring:
+ PartialEq
+ Debug
+ Clone
+ Copy
// + Copy
+ Sum<<Self as Add>::Output>
+ Sum<<Self as Mul>::Output>
{
/// C defines the coefficient type
type C: Debug + Clone;
const Q: u64;
const N: usize;
fn param(&self) -> RingParam;
fn coeffs(&self) -> Vec<Self::C>;
fn zero() -> Self;
fn zero(param: &RingParam) -> Self;
// note/wip/warning: dist (0,q) with f64, will output more '0=q' elements than other values
fn rand(rng: impl Rng, dist: impl Distribution<f64>) -> Self;
fn rand(rng: impl Rng, dist: impl Distribution<f64>, param: &RingParam) -> Self;
fn from_vec(coeffs: Vec<Self::C>) -> Self;
fn from_vec(param: &RingParam, coeffs: Vec<Self::C>) -> Self;
fn decompose(&self, beta: u32, l: u32) -> Vec<Self>;
fn remodule<const P: u64>(&self) -> impl Ring;
fn mod_switch<const P: u64>(&self) -> impl Ring;
fn remodule(&self, p:u64) -> impl Ring;
fn mod_switch(&self, p:u64) -> impl Ring;
/// returns [ [(num/den) * self].round() ] mod q
/// ie. performs the multiplication and division over f64, and then it

355
arith/src/ring_n.rs
View File

@@ -1,45 +1,43 @@
//! Polynomial ring Z[X]/(X^N+1)
//!
use anyhow::Result;
use itertools::zip_eq;
use rand::{distributions::Distribution, Rng};
use std::array;
use std::fmt;
use std::iter::Sum;
use std::ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign};
use crate::Ring;
use std::ops::{Add, AddAssign, Mul, Neg, Sub, SubAssign};
// 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> {
impl<const N: usize> R<N> {
// type C = i64;
// const Q: u64 = i64::MAX as u64; // WIP
// const N: usize = N;
#[derive(Clone)]
pub struct R {
pub n: usize,
pub coeffs: Vec<i64>,
}
impl R {
pub fn coeffs(&self) -> Vec<i64> {
self.0.to_vec()
self.coeffs.clone()
}
fn zero() -> Self {
let coeffs: [i64; N] = array::from_fn(|_| 0i64);
Self(coeffs)
fn zero(n: usize) -> Self {
Self {
n,
coeffs: vec![0i64; n],
}
}
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)
fn rand(mut rng: impl Rng, dist: impl Distribution<f64>, n: usize) -> Self {
Self {
n,
coeffs: std::iter::repeat_with(|| dist.sample(&mut rng).round() as i64)
.take(n)
.collect(),
}
}
pub fn from_vec(coeffs: Vec<i64>) -> Self {
pub fn from_vec(n: usize, coeffs: Vec<i64>) -> Self {
let mut p = coeffs;
modulus::<N>(&mut p);
Self(array::from_fn(|i| p[i]))
modulus(n, &mut p);
Self { n, coeffs: p }
}
/*
@@ -71,34 +69,38 @@ impl<const N: usize> R<N> {
*/
}
impl<const Q: u64, const N: usize> From<crate::ring_nq::Rq<Q, N>> for R<N> {
fn from(rq: crate::ring_nq::Rq<Q, N>) -> Self {
Self::from_vec_u64(rq.coeffs().to_vec().iter().map(|e| e.0).collect())
impl From<crate::ring_nq::Rq> for R {
fn from(rq: crate::ring_nq::Rq) -> Self {
Self::from_vec_u64(
rq.param.n,
rq.coeffs().to_vec().iter().map(|e| e.v).collect(),
)
}
}
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)
impl R {
pub fn to_rq(self, q: u64) -> crate::Rq {
crate::Rq::from((q, self))
}
// this method is mostly for tests
pub fn from_vec_u64(coeffs: Vec<u64>) -> Self {
pub fn from_vec_u64(n: usize, coeffs: Vec<u64>) -> Self {
let coeffs_i64 = coeffs.iter().map(|c| *c as i64).collect();
Self::from_vec(coeffs_i64)
Self::from_vec(n, coeffs_i64)
}
pub fn from_vec_f64(coeffs: Vec<f64>) -> Self {
pub fn from_vec_f64(n: usize, coeffs: Vec<f64>) -> Self {
let coeffs_i64 = coeffs.iter().map(|c| c.round() as i64).collect();
Self::from_vec(coeffs_i64)
Self::from_vec(n, coeffs_i64)
}
pub fn new(coeffs: [i64; N]) -> Self {
Self(coeffs)
pub fn new(n: usize, coeffs: Vec<i64>) -> Self {
assert_eq!(n, coeffs.len());
Self { n, coeffs }
}
pub fn mul_by_i64(&self, s: i64) -> Self {
Self(array::from_fn(|i| self.0[i] * s))
Self {
n: self.n,
coeffs: self.coeffs.iter().map(|c_i| c_i * s).collect(),
}
}
pub fn infinity_norm(&self) -> u64 {
@@ -108,10 +110,10 @@ impl<const N: usize> R<N> {
.map(|x| x.abs() as u64)
.fold(0, |a, b| a.max(b))
}
pub fn mod_centered_q<const Q: u64>(&self) -> R<N> {
let q = Q as i64;
pub fn mod_centered_q(&self, q: u64) -> R {
let q = q as i64;
let r = self
.0
.coeffs
.iter()
.map(|v| {
let mut res = v % q;
@@ -121,190 +123,216 @@ impl<const N: usize> R<N> {
res
})
.collect::<Vec<i64>>();
R::<N>::from_vec(r)
R::from_vec(self.n, r)
}
}
pub fn mul_div_round<const Q: u64, const N: usize>(
v: Vec<i64>,
num: u64,
den: u64,
) -> crate::Rq<Q, N> {
pub fn mul_div_round(q: u64, n: usize, v: Vec<i64>, num: u64, den: u64) -> crate::Rq {
// 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)
crate::Rq::from_vec_f64(&crate::ring::RingParam { q, n }, r)
}
// TODO rename to make it clear that is not mod q, but mod X^N+1
// apply mod (X^N+1)
pub fn modulus<const N: usize>(p: &mut Vec<i64>) {
if p.len() < N {
pub fn modulus(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();
for i in n..p.len() {
p[i - n] = p[i - n].clone() - p[i].clone();
p[i] = 0;
}
p.truncate(N);
p.truncate(n);
}
pub fn modulus_i128<const N: usize>(p: &mut Vec<i128>) {
if p.len() < N {
pub fn modulus_i128(n: usize, p: &mut Vec<i128>) {
if p.len() < n {
return;
}
for i in N..p.len() {
p[i - N] = p[i - N].clone() - p[i].clone();
for i in n..p.len() {
p[i - n] = p[i - n].clone() - p[i].clone();
p[i] = 0;
}
p.truncate(N);
p.truncate(n);
}
impl<const N: usize> PartialEq for R<N> {
impl PartialEq for R {
fn eq(&self, other: &Self) -> bool {
self.0 == other.0
self.coeffs == other.coeffs && self.n == other.n
}
}
impl<const N: usize> Add<R<N>> for R<N> {
impl Add<R> for R {
type Output = Self;
fn add(self, rhs: Self) -> Self {
Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
assert_eq!(self.n, rhs.n);
Self {
n: self.n,
coeffs: zip_eq(self.coeffs, rhs.coeffs)
.map(|(l, r)| l + r)
.collect(),
}
}
}
impl<const N: usize> Add<&R<N>> for &R<N> {
type Output = R<N>;
impl Add<&R> for &R {
type Output = R;
fn add(self, rhs: &R<N>) -> Self::Output {
R(array::from_fn(|i| self.0[i] + rhs.0[i]))
fn add(self, rhs: &R) -> Self::Output {
assert_eq!(self.n, rhs.n);
R {
n: self.n,
coeffs: zip_eq(self.coeffs.clone(), rhs.coeffs.clone())
.map(|(l, r)| l + r)
.collect(),
}
}
}
impl<const N: usize> AddAssign for R<N> {
impl AddAssign for R {
fn add_assign(&mut self, rhs: Self) {
for i in 0..N {
self.0[i] += rhs.0[i];
assert_eq!(self.n, rhs.n);
for i in 0..self.n {
self.coeffs[i] += rhs.coeffs[i];
}
}
}
impl<const N: usize> Sum<R<N>> for R<N> {
fn sum<I>(iter: I) -> Self
impl Sum<R> for R {
fn sum<I>(mut iter: I) -> Self
where
I: Iterator<Item = Self>,
{
let mut acc = R::<N>::zero();
for e in iter {
acc += e;
}
acc
let first = iter.next().unwrap();
iter.fold(first, |acc, x| acc + x)
}
}
impl<const N: usize> Sub<R<N>> for R<N> {
impl Sub<R> for R {
type Output = Self;
fn sub(self, rhs: Self) -> Self {
Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
assert_eq!(self.n, rhs.n);
Self {
n: self.n,
coeffs: zip_eq(self.coeffs, rhs.coeffs)
.map(|(l, r)| l - r)
.collect(),
}
}
}
impl<const N: usize> Sub<&R<N>> for &R<N> {
type Output = R<N>;
impl Sub<&R> for &R {
type Output = R;
fn sub(self, rhs: &R<N>) -> Self::Output {
R(array::from_fn(|i| self.0[i] - rhs.0[i]))
}
}
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];
fn sub(self, rhs: &R) -> Self::Output {
assert_eq!(self.n, rhs.n);
R {
n: self.n,
coeffs: zip_eq(&self.coeffs, &rhs.coeffs)
.map(|(l, r)| l - r)
.collect(),
}
}
}
impl<const N: usize> Mul<R<N>> for R<N> {
impl SubAssign for R {
fn sub_assign(&mut self, rhs: Self) {
assert_eq!(self.n, rhs.n);
for i in 0..self.n {
self.coeffs[i] -= rhs.coeffs[i];
}
}
}
impl Mul<R> for R {
type Output = Self;
fn mul(self, rhs: Self) -> Self {
naive_poly_mul(&self, &rhs)
}
}
impl<const N: usize> Mul<&R<N>> for &R<N> {
type Output = R<N>;
impl Mul<&R> for &R {
type Output = R;
fn mul(self, rhs: &R<N>) -> Self::Output {
fn mul(self, rhs: &R) -> Self::Output {
naive_poly_mul(self, rhs)
}
}
// TODO WIP
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 {
pub fn naive_poly_mul(poly1: &R, poly2: &R) -> R {
assert_eq!(poly1.n, poly2.n);
let n = poly1.n;
let poly1: Vec<i128> = poly1.coeffs.iter().map(|c| *c as i128).collect();
let poly2: Vec<i128> = poly2.coeffs.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];
}
}
// apply mod (X^N + 1))
// R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
modulus_i128::<N>(&mut result);
modulus_i128(n, &mut result);
// dbg!(&result);
// dbg!(R::<N>(array::from_fn(|i| result[i] as i64)).coeffs());
let result_i64: Vec<i64> = result.iter().map(|c_i| *c_i as i64).collect();
let r = R::from_vec(n, result_i64);
// sanity check: check that there are no coeffs > i64_max
assert_eq!(
result,
R::<N>(array::from_fn(|i| result[i] as i64))
.coeffs()
.iter()
.map(|c| *c as i128)
.collect::<Vec<_>>()
r.coeffs.iter().map(|c| *c as i128).collect::<Vec<_>>()
);
R(array::from_fn(|i| result[i] as i64))
r
}
pub fn naive_mul_2<const N: usize>(poly1: &Vec<i128>, poly2: &Vec<i128>) -> Vec<i128> {
let mut result: Vec<i128> = vec![0; (N * 2) - 1];
for i in 0..N {
for j in 0..N {
pub fn naive_mul_2(n: usize, poly1: &Vec<i128>, poly2: &Vec<i128>) -> Vec<i128> {
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];
}
}
// apply mod (X^N + 1))
// R::<N>::from_vec(result.iter().map(|c| *c as i64).collect())
modulus_i128::<N>(&mut result);
modulus_i128(n, &mut result);
result
}
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 {
pub fn naive_mul(poly1: &R, poly2: &R) -> Vec<i64> {
assert_eq!(poly1.n, poly2.n);
let n = poly1.n;
let poly1: Vec<i128> = poly1.coeffs.iter().map(|c| *c as i128).collect();
let poly2: Vec<i128> = poly2.coeffs.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()
}
pub fn naive_mul_TMP<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<i128> = vec![0; (N * 2) - 1];
for i in 0..N {
for j in 0..N {
pub fn naive_mul_TMP(poly1: &R, poly2: &R) -> Vec<i64> {
assert_eq!(poly1.n, poly2.n);
let n = poly1.n;
let poly1: Vec<i128> = poly1.coeffs.iter().map(|c| *c as i128).collect();
let poly2: Vec<i128> = poly2.coeffs.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];
}
}
// dbg!(&result);
modulus_i128::<N>(&mut result);
modulus_i128(n, &mut result);
// for c_i in result.iter() {
// println!("---");
// println!("{:?}", &c_i);
@@ -316,8 +344,8 @@ pub fn naive_mul_TMP<const N: usize>(poly1: &R<N>, poly2: &R<N>) -> Vec<i64> {
}
// wip
pub fn mod_centered_q<const Q: u64, const N: usize>(p: Vec<i128>) -> R<N> {
let q: i128 = Q as i128;
pub fn mod_centered_q(q: u64, n: usize, p: Vec<i128>) -> R {
let q: i128 = q as i128;
let r = p
.iter()
.map(|v| {
@@ -328,10 +356,10 @@ pub fn mod_centered_q<const Q: u64, const N: usize>(p: Vec<i128>) -> R<N> {
res
})
.collect::<Vec<i128>>();
R::<N>::from_vec(r.iter().map(|v| *v as i64).collect::<Vec<i64>>())
R::from_vec(n, r.iter().map(|v| *v as i64).collect::<Vec<i64>>())
}
impl<const N: usize> Mul<i64> for R<N> {
impl Mul<i64> for R {
type Output = Self;
fn mul(self, s: i64) -> Self {
@@ -339,34 +367,38 @@ impl<const N: usize> Mul<i64> for R<N> {
}
}
// mul by u64
impl<const N: usize> Mul<u64> for R<N> {
impl Mul<u64> for R {
type Output = Self;
fn mul(self, s: u64) -> Self {
self.mul_by_i64(s as i64)
}
}
impl<const N: usize> Mul<&u64> for &R<N> {
type Output = R<N>;
impl Mul<&u64> for &R {
type Output = R;
fn mul(self, s: &u64) -> Self::Output {
self.mul_by_i64(*s as i64)
}
}
impl<const N: usize> Neg for R<N> {
impl Neg for R {
type Output = Self;
fn neg(self) -> Self::Output {
Self(array::from_fn(|i| -self.0[i]))
// Self(array::from_fn(|i| -self.0[i]))
Self {
n: self.n,
coeffs: self.coeffs.iter().map(|c_i| -c_i).collect(),
}
}
}
impl<const N: usize> R<N> {
impl R {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
let mut str = "";
let mut zero = true;
for (i, coeff) in self.0.iter().enumerate().rev() {
for (i, coeff) in self.coeffs.iter().enumerate().rev() {
if *coeff == 0 {
continue;
}
@@ -395,18 +427,18 @@ impl<const N: usize> R<N> {
f.write_str(" mod Z")?;
f.write_str("/(X^")?;
f.write_str(N.to_string().as_str())?;
f.write_str(self.n.to_string().as_str())?;
f.write_str("+1)")?;
Ok(())
}
}
impl<const N: usize> fmt::Display for R<N> {
impl fmt::Display for R {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.fmt(f)?;
Ok(())
}
}
impl<const N: usize> fmt::Debug for R<N> {
impl fmt::Debug for R {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.fmt(f)?;
Ok(())
@@ -420,38 +452,33 @@ mod tests {
#[test]
fn test_mul() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 2;
let q: i64 = Q as i64;
let n: usize = 2;
let q: i64 = (2u64.pow(16) + 1) as i64;
// test vectors generated with SageMath
let a: [i64; N] = [q - 1, q - 1];
let b: [i64; N] = [q - 1, q - 1];
let c: [i64; N] = [0, 8589934592];
test_mul_opt::<Q, N>(a, b, c)?;
let a: Vec<i64> = vec![q - 1, q - 1];
let b: Vec<i64> = vec![q - 1, q - 1];
let c: Vec<i64> = vec![0, 8589934592];
test_mul_opt(n, a, b, c)?;
let a: [i64; N] = [1, q - 1];
let b: [i64; N] = [1, q - 1];
let c: [i64; N] = [-4294967295, 131072];
test_mul_opt::<Q, N>(a, b, c)?;
let a: Vec<i64> = vec![1, q - 1];
let b: Vec<i64> = vec![1, q - 1];
let c: Vec<i64> = vec![-4294967295, 131072];
test_mul_opt(n, a, b, c)?;
Ok(())
}
fn test_mul_opt<const Q: u64, const N: usize>(
a: [i64; N],
b: [i64; N],
expected_c: [i64; N],
) -> Result<()> {
let mut a = R::new(a);
let mut b = R::new(b);
fn test_mul_opt(n: usize, a: Vec<i64>, b: Vec<i64>, expected_c: Vec<i64>) -> Result<()> {
let mut a = R::new(n, a);
let mut b = R::new(n, b);
dbg!(&a);
dbg!(&b);
let expected_c = R::new(expected_c);
let expected_c = R::new(n, expected_c);
let mut c = naive_mul(&mut a, &mut b);
modulus::<N>(&mut c);
dbg!(R::<N>::from_vec(c.clone()));
assert_eq!(c, expected_c.0.to_vec());
modulus(n, &mut c);
dbg!(R::from_vec(n, c.clone()));
assert_eq!(c, expected_c.coeffs);
Ok(())
}
}

525
arith/src/ring_nq.rs
View File

@@ -2,8 +2,8 @@
//!
use anyhow::{anyhow, Result};
use itertools::zip_eq;
use rand::{distributions::Distribution, Rng};
use std::array;
use std::fmt;
use std::iter::Sum;
use std::ops::{Add, AddAssign, Mul, Neg, Sub, SubAssign};
@@ -11,82 +11,91 @@ use std::ops::{Add, AddAssign, Mul, Neg, Sub, SubAssign};
use crate::ntt::NTT;
use crate::zq::{modulus_u64, Zq};
use crate::Ring;
use crate::{Ring, RingParam};
// NOTE: currently using fixed-size arrays, but pending to see if with
// real-world parameters the stack can keep up; if not will move everything to
// use Vec.
/// 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],
#[derive(Clone)]
pub struct Rq {
pub param: RingParam,
pub(crate) coeffs: Vec<Zq>,
// 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]>,
pub(crate) evals: Option<Vec<Zq>>,
}
impl<const Q: u64, const N: usize> Ring for Rq<Q, N> {
type C = Zq<Q>;
const Q: u64 = Q;
const N: usize = N;
impl Ring for Rq {
type C = Zq;
fn param(&self) -> RingParam {
self.param
}
fn coeffs(&self) -> Vec<Self::C> {
self.coeffs.to_vec()
}
fn zero() -> Self {
let coeffs = array::from_fn(|_| Zq::zero());
fn zero(param: &RingParam) -> Self {
Self {
coeffs,
param: param.clone(),
coeffs: vec![Zq::zero(param.q); param.n],
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));
fn rand(mut rng: impl Rng, dist: impl Distribution<f64>, param: &RingParam) -> Self {
Self {
coeffs,
param: param.clone(),
coeffs: std::iter::repeat_with(|| Self::C::rand(&mut rng, &dist, param.q))
.take(param.n)
.collect(),
evals: None,
}
}
fn from_vec(coeffs: Vec<Zq<Q>>) -> Self {
fn from_vec(param: &RingParam, coeffs: Vec<Zq>) -> Self {
let mut p = coeffs;
modulus::<Q, N>(&mut p);
let coeffs = array::from_fn(|i| p[i]);
modulus(param.q, param.n, &mut p);
Self {
coeffs,
param: param.clone(),
coeffs: p,
evals: None,
}
}
// returns the decomposition of each polynomial coefficient, such
// decomposition will be a vecotor of length N, containint N vectors of Zq
// decomposition will be a vector of length N, containing N vectors of Zq
fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
let elems: Vec<Vec<Zq<Q>>> = self.coeffs.iter().map(|r| r.decompose(beta, l)).collect();
let elems: Vec<Vec<Zq>> = self.coeffs.iter().map(|r| r.decompose(beta, l)).collect();
// transpose it
let r: Vec<Vec<Zq<Q>>> = (0..elems[0].len())
let r: Vec<Vec<Zq>> = (0..elems[0].len())
.map(|i| (0..elems.len()).map(|j| elems[j][i]).collect())
.collect();
// convert it to Rq<Q,N>
r.iter().map(|a_i| Self::from_vec(a_i.clone())).collect()
r.iter()
.map(|a_i| Self::from_vec(&self.param, a_i.clone()))
.collect()
}
// 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
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())
fn remodule(&self, p: u64) -> Rq {
let param = RingParam {
q: p,
n: self.param.n,
};
Rq::from_vec_u64(&param, self.coeffs().iter().map(|m_i| m_i.v).collect())
}
/// perform the mod switch operation from Q to Q', where Q2=Q'
// fn mod_switch<const P: u64, const M: usize>(&self) -> impl Ring {
fn mod_switch<const P: u64>(&self) -> Rq<P, N> {
// assert_eq!(N, M); // sanity check
Rq::<P, N> {
coeffs: array::from_fn(|i| self.coeffs[i].mod_switch::<P>()),
fn mod_switch(&self, p: u64) -> Rq {
let param = RingParam {
q: p,
n: self.param.n,
};
Rq {
param,
coeffs: self.coeffs.iter().map(|c_i| c_i.mod_switch(p)).collect(),
evals: None,
}
}
@@ -98,105 +107,138 @@ impl<const Q: u64, const N: usize> Ring for Rq<Q, N> {
let r: Vec<f64> = self
.coeffs()
.iter()
.map(|e| ((num as f64 * e.0 as f64) / den as f64).round())
.map(|e| ((num as f64 * e.v as f64) / den as f64).round())
.collect();
Rq::<Q, N>::from_vec_f64(r)
Rq::from_vec_f64(&self.param, r)
}
}
impl<const Q: u64, const N: usize> From<crate::ring_n::R<N>> for Rq<Q, N> {
fn from(r: crate::ring_n::R<N>) -> Self {
impl From<(u64, crate::ring_n::R)> for Rq {
fn from(qr: (u64, crate::ring_n::R)) -> Self {
let (q, r) = qr;
assert_eq!(r.n, r.coeffs.len());
Self::from_vec(
&RingParam { q, n: r.n },
r.coeffs()
.iter()
.map(|e| Zq::<Q>::from_f64(*e as f64))
.map(|e| Zq::from_f64(q, *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 {
pub fn modulus(q: u64, n: usize, p: &mut Vec<Zq>) {
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);
for i in n..p.len() {
p[i - n] = p[i - n].clone() - p[i].clone();
p[i] = Zq::zero(q);
}
p.truncate(N);
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
impl Rq {
pub fn coeffs(&self) -> Vec<Zq> {
self.coeffs.clone()
}
pub fn compute_evals(&mut self) {
self.evals = Some(NTT::<Q, N>::ntt(self.coeffs));
self.evals = Some(NTT::ntt(self).coeffs);
// TODO improve, ntt returns Rq but here just needs Vec<Zq>
}
pub fn to_r(self) -> crate::R<N> {
crate::R::<N>::from(self)
pub fn to_r(self) -> crate::R {
crate::R::from(self)
}
// TODO rm since it is implemented in Ring trait impl
// pub fn zero() -> Self {
// let coeffs = array::from_fn(|_| Zq::zero());
// Self {
// coeffs,
// evals: None,
// }
// }
// 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_u64(param: &RingParam, coeffs: Vec<u64>) -> Self {
let coeffs_mod_q: Vec<Zq> = coeffs.iter().map(|c| Zq::from_u64(param.q, *c)).collect();
Self::from_vec(param, 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_f64(param: &RingParam, coeffs: Vec<f64>) -> Self {
let coeffs_mod_q: Vec<Zq> = coeffs.iter().map(|c| Zq::from_f64(param.q, *c)).collect();
Self::from_vec(param, 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 from_vec_i64(param: &RingParam, coeffs: Vec<i64>) -> Self {
let coeffs_mod_q: Vec<Zq> = coeffs
.iter()
.map(|c| Zq::from_f64(param.q, *c as f64))
.collect();
Self::from_vec(param, coeffs_mod_q)
}
pub fn new(coeffs: [Zq<Q>; N], evals: Option<[Zq<Q>; N]>) -> Self {
Self { coeffs, evals }
pub fn new(param: &RingParam, coeffs: Vec<Zq>, evals: Option<Vec<Zq>>) -> Self {
Self {
param: *param,
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()));
pub fn rand_abs(
mut rng: impl Rng,
dist: impl Distribution<f64>,
param: &RingParam,
) -> Result<Self> {
Ok(Self {
coeffs,
param: *param,
coeffs: std::iter::repeat_with(|| Zq::from_f64(param.q, dist.sample(&mut rng).abs()))
.take(param.n)
.collect(),
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()));
pub fn rand_f64_abs(
mut rng: impl Rng,
dist: impl Distribution<f64>,
param: &RingParam,
) -> Result<Self> {
Ok(Self {
coeffs,
param: *param,
coeffs: std::iter::repeat_with(|| Zq::from_f64(param.q, dist.sample(&mut rng).abs()))
.take(param.n)
.collect(),
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)));
pub fn rand_f64(
mut rng: impl Rng,
dist: impl Distribution<f64>,
param: &RingParam,
) -> Result<Self> {
Ok(Self {
coeffs,
param: *param,
coeffs: std::iter::repeat_with(|| Zq::from_f64(param.q, dist.sample(&mut rng)))
.take(param.n)
.collect(),
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)));
pub fn rand_u64(
mut rng: impl Rng,
dist: impl Distribution<u64>,
param: &RingParam,
) -> Result<Self> {
Ok(Self {
coeffs,
param: *param,
coeffs: std::iter::repeat_with(|| Zq::from_u64(param.q, dist.sample(&mut rng)))
.take(param.n)
.collect(),
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)));
pub fn rand_bin(
mut rng: impl Rng,
dist: impl Distribution<bool>,
param: &RingParam,
) -> Result<Self> {
Ok(Rq {
coeffs,
param: *param,
coeffs: std::iter::repeat_with(|| Zq::from_bool(param.q, dist.sample(&mut rng)))
.take(param.n)
.collect(),
evals: None,
})
}
@@ -208,36 +250,43 @@ impl<const Q: u64, const N: usize> Rq<Q, N> {
// }
// applies mod(T) to all coefficients of self
pub fn coeffs_mod<const T: u64>(&self) -> Self {
Rq::<Q, N>::from_vec_u64(
pub fn coeffs_mod(&self, param: &RingParam, t: u64) -> Self {
Rq::from_vec_u64(
param,
self.coeffs()
.iter()
.map(|m_i| modulus_u64::<T>(m_i.0))
.map(|m_i| modulus_u64(t, m_i.v))
.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>>> {
pub fn mul_by_matrix(&self, m: &Vec<Vec<Zq>>) -> Result<Vec<Zq>> {
matrix_vec_product(m, &self.coeffs.to_vec())
}
pub fn mul_by_zq(&self, s: &Zq<Q>) -> Self {
pub fn mul_by_zq(&self, s: &Zq) -> Self {
Self {
coeffs: array::from_fn(|i| self.coeffs[i] * *s),
param: self.param,
coeffs: self.coeffs.iter().map(|c_i| *c_i * *s).collect(),
evals: None,
}
}
pub fn mul_by_u64(&self, s: u64) -> Self {
let s = Zq::from_u64(s);
let s = Zq::from_u64(self.param.q, s);
Self {
coeffs: array::from_fn(|i| self.coeffs[i] * s),
// coeffs: self.coeffs.iter().map(|&e| e * s).collect(),
param: self.param,
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)),
param: self.param,
coeffs: self
.coeffs
.iter()
.map(|c_i| Zq::from_f64(self.param.q, c_i.v as f64 * s))
.collect(),
evals: None,
}
}
@@ -251,9 +300,9 @@ impl<const Q: u64, const N: usize> Rq<Q, N> {
let r: Vec<f64> = self
.coeffs()
.iter()
.map(|e| (e.0 as f64 / s as f64).round())
.map(|e| (e.v as f64 / s as f64).round())
.collect();
Rq::<Q, N>::from_vec_f64(r)
Rq::from_vec_f64(&self.param, r)
}
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
@@ -261,19 +310,19 @@ impl<const Q: u64, const N: usize> Rq<Q, N> {
let mut str = "";
let mut zero = true;
for (i, coeff) in self.coeffs.iter().enumerate().rev() {
if coeff.0 == 0 {
if coeff.v == 0 {
continue;
}
zero = false;
f.write_str(str)?;
if coeff.0 != 1 {
f.write_str(coeff.0.to_string().as_str())?;
if coeff.v != 1 {
f.write_str(coeff.v.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 coeff.v == 1 && i == 0 {
f.write_str(coeff.v.to_string().as_str())?;
}
if i == 1 {
f.write_str("x")?;
@@ -288,9 +337,9 @@ impl<const Q: u64, const N: usize> Rq<Q, N> {
}
f.write_str(" mod Z_")?;
f.write_str(Q.to_string().as_str())?;
f.write_str(self.param.q.to_string().as_str())?;
f.write_str("/(X^")?;
f.write_str(N.to_string().as_str())?;
f.write_str(self.param.n.to_string().as_str())?;
f.write_str("+1)")?;
Ok(())
}
@@ -298,16 +347,20 @@ impl<const Q: u64, const N: usize> Rq<Q, N> {
pub fn infinity_norm(&self) -> u64 {
self.coeffs()
.iter()
.map(|x| if x.0 > (Q / 2) { Q - x.0 } else { x.0 })
.map(|x| {
if x.v > (self.param.q / 2) {
self.param.q - x.v
} else {
x.v
}
})
.fold(0, |a, b| a.max(b))
}
pub fn mod_centered_q(&self) -> crate::ring_n::R<N> {
self.to_r().mod_centered_q::<Q>()
pub fn mod_centered_q(&self) -> crate::ring_n::R {
self.clone().to_r().mod_centered_q(self.param.q)
}
}
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());
pub fn matrix_vec_product(m: &Vec<Vec<Zq>>, v: &Vec<Zq>) -> Result<Vec<Zq>> {
if m.len() != m[0].len() {
return Err(anyhow!("expected 'm' to be a square matrix"));
}
@@ -319,6 +372,8 @@ pub fn matrix_vec_product<const Q: u64>(m: &Vec<Vec<Zq<Q>>>, v: &Vec<Zq<Q>>) ->
));
}
assert_eq!(m[0][0].q, v[0].q); // TODO change to returning err
Ok(m.iter()
.map(|row| {
row.iter()
@@ -326,12 +381,15 @@ pub fn matrix_vec_product<const Q: u64>(m: &Vec<Vec<Zq<Q>>>, v: &Vec<Zq<Q>>) ->
.map(|(&row_i, &v_i)| row_i * v_i)
.sum()
})
.collect::<Vec<Zq<Q>>>())
.collect::<Vec<Zq>>())
}
pub fn transpose<const Q: u64>(m: &[Vec<Zq<Q>>]) -> Vec<Vec<Zq<Q>>> {
pub fn transpose(m: &[Vec<Zq>]) -> Vec<Vec<Zq>> {
assert!(m.len() > 0);
assert!(m[0].len() > 0);
let q = m[0][0].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()];
let mut r: Vec<Vec<Zq>> = vec![vec![Zq::zero(q); 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;
@@ -340,205 +398,221 @@ pub fn transpose<const Q: u64>(m: &[Vec<Zq<Q>>]) -> Vec<Vec<Zq<Q>>> {
r
}
impl<const Q: u64, const N: usize> PartialEq for Rq<Q, N> {
impl PartialEq for Rq {
fn eq(&self, other: &Self) -> bool {
self.coeffs == other.coeffs
self.coeffs == other.coeffs && self.param == other.param
}
}
impl<const Q: u64, const N: usize> Add<Rq<Q, N>> for Rq<Q, N> {
impl Add<Rq> for Rq {
type Output = Self;
fn add(self, rhs: Self) -> Self {
assert_eq!(self.param, rhs.param);
Self {
coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
param: self.param,
coeffs: zip_eq(self.coeffs, rhs.coeffs)
.map(|(l, r)| l + r)
.collect(),
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> Add<&Rq<Q, N>> for &Rq<Q, N> {
type Output = Rq<Q, N>;
impl Add<&Rq> for &Rq {
type Output = Rq;
fn add(self, rhs: &Rq<Q, N>) -> Self::Output {
fn add(self, rhs: &Rq) -> Self::Output {
assert_eq!(self.param, rhs.param);
Rq {
coeffs: array::from_fn(|i| self.coeffs[i] + rhs.coeffs[i]),
param: self.param,
coeffs: zip_eq(self.coeffs.clone(), rhs.coeffs.clone())
.map(|(l, r)| l + r)
.collect(),
evals: None,
}
}
}
impl<const Q: u64, const N: usize> AddAssign for Rq<Q, N> {
impl AddAssign for Rq {
fn add_assign(&mut self, rhs: Self) {
for i in 0..N {
debug_assert_eq!(self.param, rhs.param);
for i in 0..self.param.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
impl Sum<Rq> for Rq {
fn sum<I>(mut iter: I) -> Self
where
I: Iterator<Item = Self>,
{
let mut acc = Rq::<Q, N>::zero();
for e in iter {
acc += e;
}
acc
let first = iter.next().unwrap();
iter.fold(first, |acc, x| acc + x)
}
}
impl<const Q: u64, const N: usize> Sub<Rq<Q, N>> for Rq<Q, N> {
impl Sub<Rq> for Rq {
type Output = Self;
fn sub(self, rhs: Self) -> Self {
assert_eq!(self.param, rhs.param);
Self {
coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
param: self.param,
coeffs: zip_eq(self.coeffs, rhs.coeffs)
.map(|(l, r)| l - r)
.collect(),
evals: None,
}
}
}
impl<const Q: u64, const N: usize> Sub<&Rq<Q, N>> for &Rq<Q, N> {
type Output = Rq<Q, N>;
impl Sub<&Rq> for &Rq {
type Output = Rq;
fn sub(self, rhs: &Rq<Q, N>) -> Self::Output {
fn sub(self, rhs: &Rq) -> Self::Output {
debug_assert_eq!(self.param, rhs.param);
Rq {
coeffs: array::from_fn(|i| self.coeffs[i] - rhs.coeffs[i]),
param: self.param,
coeffs: zip_eq(self.coeffs.clone(), rhs.coeffs.clone())
.map(|(l, r)| l - r)
.collect(),
evals: None,
}
}
}
impl<const Q: u64, const N: usize> SubAssign for Rq<Q, N> {
impl SubAssign for Rq {
fn sub_assign(&mut self, rhs: Self) {
for i in 0..N {
debug_assert_eq!(self.param, rhs.param);
for i in 0..self.param.n {
self.coeffs[i] -= rhs.coeffs[i];
}
}
}
impl<const Q: u64, const N: usize> Mul<Rq<Q, N>> for Rq<Q, N> {
impl Mul<Rq> for Rq {
type Output = Self;
fn mul(self, rhs: Self) -> Self {
mul(&self, &rhs)
}
}
impl<const Q: u64, const N: usize> Mul<&Rq<Q, N>> for &Rq<Q, N> {
type Output = Rq<Q, N>;
impl Mul<&Rq> for &Rq {
type Output = Rq;
fn mul(self, rhs: &Rq<Q, N>) -> Self::Output {
fn mul(self, rhs: &Rq) -> Self::Output {
mul(self, rhs)
}
}
// mul by Zq element
impl<const Q: u64, const N: usize> Mul<Zq<Q>> for Rq<Q, N> {
impl Mul<Zq> for Rq {
type Output = Self;
fn mul(self, s: Zq<Q>) -> Self {
fn mul(self, s: Zq) -> Self {
self.mul_by_zq(&s)
}
}
impl<const Q: u64, const N: usize> Mul<&Zq<Q>> for &Rq<Q, N> {
type Output = Rq<Q, N>;
impl Mul<&Zq> for &Rq {
type Output = Rq;
fn mul(self, s: &Zq<Q>) -> Self::Output {
fn mul(self, s: &Zq) -> Self::Output {
self.mul_by_zq(s)
}
}
// mul by u64
impl<const Q: u64, const N: usize> Mul<u64> for Rq<Q, N> {
impl Mul<u64> for Rq {
type Output = Self;
fn mul(self, s: u64) -> Self {
self.mul_by_u64(s)
}
}
impl<const Q: u64, const N: usize> Mul<&u64> for &Rq<Q, N> {
type Output = Rq<Q, N>;
impl Mul<&u64> for &Rq {
type Output = Rq;
fn mul(self, s: &u64) -> Self::Output {
self.mul_by_u64(*s)
}
}
// mul by f64
impl<const Q: u64, const N: usize> Mul<f64> for Rq<Q, N> {
impl Mul<f64> for Rq {
type Output = Self;
fn mul(self, s: f64) -> Self {
self.mul_by_f64(s)
}
}
impl<const Q: u64, const N: usize> Mul<&f64> for &Rq<Q, N> {
type Output = Rq<Q, N>;
impl Mul<&f64> for &Rq {
type Output = Rq;
fn mul(self, s: &f64) -> Self::Output {
self.mul_by_f64(*s)
}
}
impl<const Q: u64, const N: usize> Neg for Rq<Q, N> {
impl Neg for Rq {
type Output = Self;
fn neg(self) -> Self::Output {
Self {
coeffs: array::from_fn(|i| -self.coeffs[i]),
param: self.param,
coeffs: self.coeffs.iter().map(|c_i| -*c_i).collect(),
evals: None,
}
}
}
// note: this assumes that Q is prime
fn mul_mut<const Q: u64, const N: usize>(lhs: &mut Rq<Q, N>, rhs: &mut Rq<Q, N>) -> Rq<Q, N> {
fn mul_mut(lhs: &mut Rq, rhs: &mut Rq) -> Rq {
assert_eq!(lhs.param, rhs.param);
// reuse evaluations if already computed
if !lhs.evals.is_some() {
lhs.evals = Some(NTT::<Q, N>::ntt(lhs.coeffs));
lhs.evals = Some(NTT::ntt(lhs).coeffs);
};
if !rhs.evals.is_some() {
rhs.evals = Some(NTT::<Q, N>::ntt(rhs.coeffs));
rhs.evals = Some(NTT::ntt(rhs).coeffs);
};
let lhs_evals = lhs.evals.unwrap();
let rhs_evals = rhs.evals.unwrap();
let lhs_evals = lhs.evals.clone().unwrap();
let rhs_evals = rhs.evals.clone().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))
let c_ntt: Rq = Rq::from_vec(
&lhs.param,
zip_eq(lhs_evals, rhs_evals).map(|(l, r)| l * r).collect(),
);
let c = NTT::intt(&c_ntt);
Rq::new(&lhs.param, c.coeffs, Some(c_ntt.coeffs))
}
// note: this assumes that Q is prime
// TODO impl karatsuba for non-prime Q
fn mul<const Q: u64, const N: usize>(lhs: &Rq<Q, N>, rhs: &Rq<Q, N>) -> Rq<Q, N> {
// TODO impl karatsuba for non-prime Q. Alternatively check NTT with RNS trick.
fn mul(lhs: &Rq, rhs: &Rq) -> Rq {
assert_eq!(lhs.param, rhs.param);
// reuse evaluations if already computed
let lhs_evals = if lhs.evals.is_some() {
lhs.evals.unwrap()
let lhs_evals: Vec<Zq> = if lhs.evals.is_some() {
lhs.evals.clone().unwrap()
} else {
NTT::<Q, N>::ntt(lhs.coeffs)
NTT::ntt(lhs).coeffs
};
let rhs_evals = if rhs.evals.is_some() {
rhs.evals.unwrap()
let rhs_evals: Vec<Zq> = if rhs.evals.is_some() {
rhs.evals.clone().unwrap()
} else {
NTT::<Q, N>::ntt(rhs.coeffs)
NTT::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))
let c_ntt: Rq = Rq::from_vec(
&lhs.param,
zip_eq(lhs_evals, rhs_evals).map(|(l, r)| l * r).collect(),
);
let c = NTT::intt(&c_ntt);
Rq::new(&lhs.param, c.coeffs, Some(c_ntt.coeffs))
}
impl<const Q: u64, const N: usize> fmt::Display for Rq<Q, N> {
impl fmt::Display for Rq {
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> {
impl fmt::Debug for Rq {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.fmt(f)?;
Ok(())
@@ -552,31 +626,30 @@ mod tests {
#[test]
fn test_polynomial_ring() {
// the test values used are generated with SageMath
const Q: u64 = 7;
const N: usize = 3;
let param = RingParam { q: 7, n: 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]);
let p = Rq::from_vec_u64(&param, 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]);
let p = Rq::from_vec_u64(&param, vec![0u64, 1, param.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]);
let p = Rq::from_vec_u64(&RingParam { q: 7, n: 4 }, 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]);
let p = Rq::from_vec_u64(&param, 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]);
let p = Rq::from_vec_u64(&param, 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]);
let a = Rq::from_vec_u64(&param, 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]);
let b = Rq::from_vec_u64(&param, 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
@@ -593,34 +666,37 @@ mod tests {
#[test]
fn test_mul() -> Result<()> {
const Q: u64 = 2u64.pow(16) + 1;
const N: usize = 4;
let param = RingParam {
q: 2u64.pow(16) + 1,
n: 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: Vec<u64> = vec![1u64, 2, 3, 4];
let b: Vec<u64> = vec![1u64, 2, 3, 4];
let c: Vec<u64> = vec![65513, 65517, 65531, 20];
test_mul_opt(&param, 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)?;
let a: Vec<u64> = vec![0u64, 0, 0, 2];
let b: Vec<u64> = vec![0u64, 0, 0, 2];
let c: Vec<u64> = vec![0u64, 0, 65533, 0];
test_mul_opt(&param, a, b, c)?;
// TODO more testvectors
Ok(())
}
fn test_mul_opt<const Q: u64, const N: usize>(
a: [u64; N],
b: [u64; N],
expected_c: [u64; N],
fn test_mul_opt(
param: &RingParam,
a: Vec<u64>,
b: Vec<u64>,
expected_c: Vec<u64>,
) -> 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);
assert_eq!(a.len(), param.n);
assert_eq!(b.len(), param.n);
let mut a = Rq::from_vec_u64(&param, a);
let mut b = Rq::from_vec_u64(&param, b);
let expected_c = Rq::from_vec_u64(&param, expected_c);
let c = mul_mut(&mut a, &mut b);
assert_eq!(c, expected_c);
@@ -629,26 +705,25 @@ mod tests {
#[test]
fn test_rq_decompose() -> Result<()> {
const Q: u64 = 16;
const N: usize = 4;
let param = RingParam { q: 16, n: 4 };
let beta = 4;
let l = 2;
let a = Rq::<Q, N>::from_vec_u64(vec![7u64, 14, 3, 6]);
let a = Rq::from_vec_u64(&param, vec![7u64, 14, 3, 6]);
let d = a.decompose(beta, l);
assert_eq!(
d[0].coeffs().to_vec(),
d[0].coeffs(),
vec![1u64, 3, 0, 1]
.iter()
.map(|e| Zq::<Q>::from_u64(*e))
.map(|e| Zq::from_u64(param.q, *e))
.collect::<Vec<_>>()
);
assert_eq!(
d[1].coeffs().to_vec(),
d[1].coeffs(),
vec![3u64, 2, 3, 2]
.iter()
.map(|e| Zq::<Q>::from_u64(*e))
.map(|e| Zq::from_u64(param.q, *e))
.collect::<Vec<_>>()
);
Ok(())

281
arith/src/ring_torus.rs
View File

@@ -7,63 +7,94 @@
//! u64, we fit it into the `Ring` trait (from ring.rs) so that we can compose
//! the 𝕋_<N,q> implementation with the other objects from the code.
use itertools::zip_eq;
use rand::{distributions::Distribution, Rng};
use std::array;
use std::iter::Sum;
use std::ops::{Add, AddAssign, Mul, Neg, Sub, SubAssign};
use crate::{ring::Ring, torus::T64, Rq, Zq};
use crate::{
ring::{Ring, RingParam},
torus::T64,
Rq, Zq,
};
/// 𝕋_<N,Q>[X] = 𝕋<Q>[X]/(X^N +1), polynomials modulo X^N+1 with coefficients in
/// 𝕋, where Q=2^64.
#[derive(Clone, Copy, Debug)]
pub struct Tn<const N: usize>(pub [T64; N]);
#[derive(Clone, Debug)]
pub struct Tn {
pub param: RingParam,
pub coeffs: Vec<T64>,
}
impl<const N: usize> Ring for Tn<N> {
impl Ring for Tn {
type C = T64;
const Q: u64 = u64::MAX; // WIP
const N: usize = N;
fn param(&self) -> RingParam {
RingParam {
q: u64::MAX,
n: self.param.n,
}
}
fn coeffs(&self) -> Vec<T64> {
self.0.to_vec()
self.coeffs.to_vec()
}
fn zero() -> Self {
Self(array::from_fn(|_| T64::zero()))
fn zero(param: &RingParam) -> Self {
Self {
param: *param,
coeffs: vec![T64::zero(param); param.n],
}
}
fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
Self(array::from_fn(|_| T64::rand(&mut rng, &dist)))
fn rand(mut rng: impl Rng, dist: impl Distribution<f64>, param: &RingParam) -> Self {
Self {
param: *param,
coeffs: std::iter::repeat_with(|| T64::rand(&mut rng, &dist, &param))
.take(param.n)
.collect(),
}
}
fn from_vec(coeffs: Vec<Self::C>) -> Self {
fn from_vec(param: &RingParam, coeffs: Vec<Self::C>) -> Self {
let mut p = coeffs;
modulus::<N>(&mut p);
Self(array::from_fn(|i| p[i]))
modulus(param, &mut p);
Self {
param: *param,
coeffs: p,
}
}
fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
let elems: Vec<Vec<T64>> = self.0.iter().map(|r| r.decompose(beta, l)).collect();
let elems: Vec<Vec<T64>> = self.coeffs.iter().map(|r| r.decompose(beta, l)).collect();
// transpose it
let r: Vec<Vec<T64>> = (0..elems[0].len())
.map(|i| (0..elems.len()).map(|j| elems[j][i]).collect())
.collect();
// convert it to Tn<N>
r.iter().map(|a_i| Self::from_vec(a_i.clone())).collect()
// convert it to Tn
r.iter()
.map(|a_i| Self::from_vec(&self.param, a_i.clone()))
.collect()
}
fn remodule<const P: u64>(&self) -> Tn<N> {
fn remodule(&self, p: u64) -> Tn {
todo!()
// Rq::<P, N>::from_vec_u64(self.coeffs().iter().map(|m_i| m_i.0).collect())
}
// fn mod_switch<const P: u64>(&self) -> impl Ring {
fn mod_switch<const P: u64>(&self) -> Rq<P, N> {
fn mod_switch(&self, p: u64) -> Rq {
// unimplemented!()
// TODO WIP
let coeffs = array::from_fn(|i| Zq::<P>::from_u64(self.0[i].mod_switch::<P>().0));
Rq::<P, N> {
let coeffs = self
.coeffs
.iter()
.map(|c_i| Zq::from_u64(p, c_i.mod_switch(p).0))
.collect();
Rq {
param: RingParam {
q: p,
n: self.param.n,
},
coeffs,
evals: None,
}
@@ -78,175 +109,220 @@ impl<const N: usize> Ring for Tn<N> {
.iter()
.map(|e| T64(((num as f64 * e.0 as f64) / den as f64).round() as u64))
.collect();
Self::from_vec(r)
Self::from_vec(&self.param, r)
}
}
impl<const N: usize> Tn<N> {
impl Tn {
// multiply self by X^-h
pub fn left_rotate(&self, h: usize) -> Self {
let h = h % N;
assert!(h < N);
let c = self.0;
let n = self.param.n;
let h = h % n;
assert!(h < n);
let c = &self.coeffs;
// c[h], c[h+1], c[h+2], ..., c[n-1], -c[0], -c[1], ..., -c[h-1]
// let r: Vec<T64> = vec![c[h..N], c[0..h].iter().map(|&c_i| -c_i).collect()].concat();
let r: Vec<T64> = c[h..N]
let r: Vec<T64> = c[h..n]
.iter()
.copied()
.chain(c[0..h].iter().map(|&x| -x))
.collect();
Self::from_vec(r)
Self::from_vec(&self.param, r)
}
pub fn from_vec_u64(v: Vec<u64>) -> Self {
pub fn from_vec_u64(param: &RingParam, v: Vec<u64>) -> Self {
let coeffs = v.iter().map(|c| T64(*c)).collect();
Self::from_vec(coeffs)
Self::from_vec(param, coeffs)
}
}
// apply mod (X^N+1)
pub fn modulus<const N: usize>(p: &mut Vec<T64>) {
if p.len() < N {
pub fn modulus(param: &RingParam, p: &mut Vec<T64>) {
let n = param.n;
if p.len() < n {
return;
}
for i in N..p.len() {
p[i - N] = p[i - N].clone() - p[i].clone();
p[i] = T64::zero();
for i in n..p.len() {
p[i - n] = p[i - n].clone() - p[i].clone();
p[i] = T64::zero(param);
}
p.truncate(N);
p.truncate(n);
}
impl<const N: usize> Add<Tn<N>> for Tn<N> {
impl Add<Tn> for Tn {
type Output = Self;
fn add(self, rhs: Self) -> Self {
Self(array::from_fn(|i| self.0[i] + rhs.0[i]))
assert_eq!(self.param, rhs.param);
Self {
param: self.param,
coeffs: zip_eq(self.coeffs, rhs.coeffs)
.map(|(l, r)| l + r)
.collect(),
}
}
}
impl<const N: usize> Add<&Tn<N>> for &Tn<N> {
type Output = Tn<N>;
impl Add<&Tn> for &Tn {
type Output = Tn;
fn add(self, rhs: &Tn<N>) -> Self::Output {
Tn(array::from_fn(|i| self.0[i] + rhs.0[i]))
fn add(self, rhs: &Tn) -> Self::Output {
assert_eq!(self.param, rhs.param);
Tn {
param: self.param,
coeffs: zip_eq(self.coeffs.clone(), rhs.coeffs.clone())
.map(|(l, r)| l + r)
.collect(),
}
}
}
impl<const N: usize> AddAssign for Tn<N> {
impl AddAssign for Tn {
fn add_assign(&mut self, rhs: Self) {
for i in 0..N {
self.0[i] += rhs.0[i];
assert_eq!(self.param, rhs.param);
for i in 0..self.param.n {
self.coeffs[i] += rhs.coeffs[i];
}
}
}
impl<const N: usize> Sum<Tn<N>> for Tn<N> {
fn sum<I>(iter: I) -> Self
impl Sum<Tn> for Tn {
fn sum<I>(mut iter: I) -> Self
where
I: Iterator<Item = Self>,
{
let mut acc = Tn::<N>::zero();
for e in iter {
acc += e;
}
acc
let first = iter.next().unwrap();
iter.fold(first, |acc, x| acc + x)
}
}
impl<const N: usize> Sub<Tn<N>> for Tn<N> {
impl Sub<Tn> for Tn {
type Output = Self;
fn sub(self, rhs: Self) -> Self {
Self(array::from_fn(|i| self.0[i] - rhs.0[i]))
assert_eq!(self.param, rhs.param);
Self {
param: self.param,
coeffs: zip_eq(self.coeffs, rhs.coeffs)
.map(|(l, r)| l - r)
.collect(),
}
}
}
impl<const N: usize> Sub<&Tn<N>> for &Tn<N> {
type Output = Tn<N>;
impl Sub<&Tn> for &Tn {
type Output = Tn;
fn sub(self, rhs: &Tn<N>) -> Self::Output {
Tn(array::from_fn(|i| self.0[i] - rhs.0[i]))
}
}
impl<const N: usize> SubAssign for Tn<N> {
fn sub_assign(&mut self, rhs: Self) {
for i in 0..N {
self.0[i] -= rhs.0[i];
fn sub(self, rhs: &Tn) -> Self::Output {
assert_eq!(self.param, rhs.param);
Tn {
param: self.param,
coeffs: zip_eq(self.coeffs.clone(), rhs.coeffs.clone())
.map(|(l, r)| l - r)
.collect(),
}
}
}
impl<const N: usize> Neg for Tn<N> {
impl SubAssign for Tn {
fn sub_assign(&mut self, rhs: Self) {
assert_eq!(self.param, rhs.param);
for i in 0..self.param.n {
self.coeffs[i] -= rhs.coeffs[i];
}
}
}
impl Neg for Tn {
type Output = Self;
fn neg(self) -> Self::Output {
Tn(array::from_fn(|i| -self.0[i]))
Self {
param: self.param,
coeffs: self.coeffs.iter().map(|c_i| -*c_i).collect(),
}
}
}
impl<const N: usize> PartialEq for Tn<N> {
impl PartialEq for Tn {
fn eq(&self, other: &Self) -> bool {
self.0 == other.0
self.coeffs == other.coeffs && self.param == other.param
}
}
impl<const N: usize> Mul<Tn<N>> for Tn<N> {
impl Mul<Tn> for Tn {
type Output = Self;
fn mul(self, rhs: Self) -> Self {
naive_poly_mul(&self, &rhs)
}
}
impl<const N: usize> Mul<&Tn<N>> for &Tn<N> {
type Output = Tn<N>;
impl Mul<&Tn> for &Tn {
type Output = Tn;
fn mul(self, rhs: &Tn<N>) -> Self::Output {
fn mul(self, rhs: &Tn) -> Self::Output {
naive_poly_mul(self, rhs)
}
}
fn naive_poly_mul<const N: usize>(poly1: &Tn<N>, poly2: &Tn<N>) -> Tn<N> {
let poly1: Vec<u128> = poly1.0.iter().map(|c| c.0 as u128).collect();
let poly2: Vec<u128> = poly2.0.iter().map(|c| c.0 as u128).collect();
let mut result: Vec<u128> = vec![0; (N * 2) - 1];
for i in 0..N {
for j in 0..N {
fn naive_poly_mul(poly1: &Tn, poly2: &Tn) -> Tn {
assert_eq!(poly1.param, poly2.param);
let n = poly1.param.n;
let param = poly1.param;
let poly1: Vec<u128> = poly1.coeffs.iter().map(|c| c.0 as u128).collect();
let poly2: Vec<u128> = poly2.coeffs.iter().map(|c| c.0 as u128).collect();
let mut result: Vec<u128> = 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];
}
}
// apply mod (X^N + 1))
modulus_u128::<N>(&mut result);
// apply mod (X^n + 1))
modulus_u128(n, &mut result);
Tn(array::from_fn(|i| T64(result[i] as u64)))
Tn {
param,
coeffs: result.iter().map(|r_i| T64(*r_i as u64)).collect(),
}
}
fn modulus_u128<const N: usize>(p: &mut Vec<u128>) {
if p.len() < N {
fn modulus_u128(n: usize, p: &mut Vec<u128>) {
if p.len() < n {
return;
}
for i in N..p.len() {
// p[i - N] = p[i - N].clone() - p[i].clone();
p[i - N] = p[i - N].wrapping_sub(p[i]);
for i in n..p.len() {
// p[i - n] = p[i - n].clone() - p[i].clone();
p[i - n] = p[i - n].wrapping_sub(p[i]);
p[i] = 0;
}
p.truncate(N);
p.truncate(n);
}
impl<const N: usize> Mul<T64> for Tn<N> {
impl Mul<T64> for Tn {
type Output = Self;
fn mul(self, s: T64) -> Self {
Self(array::from_fn(|i| self.0[i] * s))
Self {
param: self.param,
coeffs: self.coeffs.iter().map(|c_i| *c_i * s).collect(),
}
}
}
// mul by u64
impl<const N: usize> Mul<u64> for Tn<N> {
impl Mul<u64> for Tn {
type Output = Self;
fn mul(self, s: u64) -> Self {
Self(array::from_fn(|i| self.0[i] * s))
Tn {
param: self.param,
coeffs: self.coeffs.iter().map(|c_i| *c_i * s).collect(),
}
}
}
impl<const N: usize> Mul<&u64> for &Tn<N> {
type Output = Tn<N>;
impl Mul<&u64> for &Tn {
type Output = Tn;
fn mul(self, s: &u64) -> Self::Output {
Tn::<N>(array::from_fn(|i| self.0[i] * *s))
Tn {
param: self.param,
coeffs: self.coeffs.iter().map(|c_i| c_i * s).collect(),
}
}
}
@@ -256,8 +332,9 @@ mod tests {
#[test]
fn test_left_rotate() {
const N: usize = 4;
let f = Tn::<N>::from_vec(
let param = RingParam { q: u64::MAX, n: 4 };
let f = Tn::from_vec(
&param,
vec![2i64, 3, -4, -1]
.iter()
.map(|c| T64(*c as u64))
@@ -267,7 +344,8 @@ mod tests {
// expect f*x^-3 == -1 -2x -3x^2 +4x^3
assert_eq!(
f.left_rotate(3),
Tn::<N>::from_vec(
Tn::from_vec(
&param,
vec![-1i64, -2, -3, 4]
.iter()
.map(|c| T64(*c as u64))
@@ -277,7 +355,8 @@ mod tests {
// expect f*x^-1 == 3 -4x -1x^2 -2x^3
assert_eq!(
f.left_rotate(1),
Tn::<N>::from_vec(
Tn::from_vec(
&param,
vec![3i64, -4, -1, -2]
.iter()
.map(|c| T64(*c as u64))

38
arith/src/torus.rs
View File

@@ -4,7 +4,7 @@ use std::{
ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
};
use crate::ring::Ring;
use crate::ring::{Ring, RingParam};
/// Let 𝕋 = ℝ/ℤ, where 𝕋 is a ℤ-module, with homogeneous external product.
/// Let 𝕋q
@@ -16,20 +16,24 @@ pub struct T64(pub u64);
// `Tn<1>`.
impl Ring for T64 {
type C = T64;
const Q: u64 = u64::MAX; // WIP
const N: usize = 1;
fn param(&self) -> RingParam {
RingParam {
q: u64::MAX, // WIP
n: 1,
}
}
fn coeffs(&self) -> Vec<T64> {
vec![self.clone()]
}
fn zero() -> Self {
fn zero(_: &RingParam) -> Self {
Self(0u64)
}
fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
fn rand(mut rng: impl Rng, dist: impl Distribution<f64>, _: &RingParam) -> Self {
let r: f64 = dist.sample(&mut rng);
Self(r.round() as u64)
}
fn from_vec(coeffs: Vec<Self::C>) -> Self {
fn from_vec(_n: &RingParam, coeffs: Vec<Self::C>) -> Self {
assert_eq!(coeffs.len(), 1);
coeffs[0]
}
@@ -37,27 +41,27 @@ impl Ring for T64 {
// TODO rm beta & l from inputs, make it always beta=2,l=64.
/// Note: only beta=2 and l=64 is supported.
fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
assert_eq!(beta, 2u32); // only beta=2 supported
// assert_eq!(l, 64u32); // only l=64 supported
assert_eq!(beta, 2u32, "only beta=2 supported");
// assert_eq!(l, 64u32, "only l=64 supported");
// (0..64)
(0..l)
(0..l as u64)
.rev()
.map(|i| T64(((self.0 >> i) & 1) as u64))
.map(|i| T64((self.0 >> i) & 1))
.collect()
}
fn remodule<const P: u64>(&self) -> T64 {
fn remodule(&self, p: u64) -> T64 {
todo!()
}
// modulus switch from Q to Q2: self * Q2/Q
fn mod_switch<const Q2: u64>(&self) -> T64 {
fn mod_switch(&self, q2: u64) -> T64 {
// for the moment we assume Q|Q2, since Q=2^64, check that Q2 is a power
// of two:
assert!(Q2.is_power_of_two());
assert!(q2.is_power_of_two());
// since Q=2^64, dividing Q2/Q is equivalent to dividing 2^log2(Q2)/2^64
// which would be like right-shifting 64-log2(Q2).
let log2_q2 = 63 - Q2.leading_zeros();
let log2_q2 = 63 - q2.leading_zeros();
T64(self.0 >> (64 - log2_q2))
}
@@ -173,9 +177,13 @@ mod tests {
let d = x.decompose(beta, l);
assert_eq!(recompose(d), T64(u64::MAX - 1));
let param = RingParam {
q: u64::MAX, // WIP
n: 1,
};
let mut rng = rand::thread_rng();
for _ in 0..1000 {
let x = T64::rand(&mut rng, Standard);
let x = T64::rand(&mut rng, Standard, &param);
let d = x.decompose(beta, l);
assert_eq!(recompose(d), x);
}

129
arith/src/tuple_ring.rs
View File

@@ -1,40 +1,46 @@
//! This file implements the struct for an Tuple of Ring Rq elements and its
//! operations, which are performed element-wise.
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::{Add, Mul, Neg, Sub},
};
use std::ops::{Add, Mul, Neg, Sub};
use crate::Ring;
use crate::{Ring, RingParam};
/// 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>(pub Vec<R>);
pub struct TR<R: Ring> {
pub k: usize,
pub r: Vec<R>,
}
// TODO rm pub from Vec<R>, so that TR can not be created from a Vec with
// invalid length, since it has to be created using the `new` method.
impl<R: Ring, const K: usize> TR<R, K> {
pub fn new(v: Vec<R>) -> Self {
assert_eq!(v.len(), K);
Self(v)
impl<R: Ring> TR<R> {
pub fn new(k: usize, r: Vec<R>) -> Self {
assert_eq!(r.len(), k);
Self { k, r }
}
pub fn zero() -> Self {
Self((0..K).into_iter().map(|_| R::zero()).collect())
pub fn zero(k: usize, r_param: &RingParam) -> Self {
Self {
k,
r: (0..k).into_iter().map(|_| R::zero(r_param)).collect(),
}
}
pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
Self(
(0..K)
pub fn rand(
mut rng: impl Rng,
dist: impl Distribution<f64>,
k: usize,
r_param: &RingParam,
) -> Self {
Self {
k,
r: (0..k)
.into_iter()
.map(|_| R::rand(&mut rng, &dist))
.map(|_| R::rand(&mut rng, &dist, r_param))
.collect(),
)
}
}
// returns the decomposition of each polynomial element
pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
@@ -43,64 +49,85 @@ impl<R: Ring, const K: usize> TR<R, K> {
}
}
impl<const K: usize> TR<crate::torus::T64, K> {
pub fn mod_switch<const Q2: u64>(&self) -> TR<crate::torus::T64, K> {
TR(self.0.iter().map(|c_i| c_i.mod_switch::<Q2>()).collect())
impl TR<crate::torus::T64> {
pub fn mod_switch(&self, q2: u64) -> TR<crate::torus::T64> {
TR::<crate::torus::T64> {
k: self.k,
r: self.r.iter().map(|c_i| c_i.mod_switch(q2)).collect(),
}
}
// pub fn mod_switch(&self, Q2: u64) -> TR<crate::torus::T64, K> {
// TR(self.0.iter().map(|c_i| c_i.mod_switch(Q2)).collect())
// }
}
impl<const N: usize, const K: usize> TR<crate::ring_torus::Tn<N>, K> {
impl TR<crate::ring_torus::Tn> {
pub fn left_rotate(&self, h: usize) -> Self {
TR(self.0.iter().map(|c_i| c_i.left_rotate(h)).collect())
TR {
k: self.k,
r: self.r.iter().map(|c_i| c_i.left_rotate(h)).collect(),
}
}
}
impl<R: Ring, const K: usize> TR<R, K> {
impl<R: Ring> TR<R> {
pub fn iter(&self) -> std::slice::Iter<R> {
self.0.iter()
self.r.iter()
}
}
impl<R: Ring, const K: usize> Add<TR<R, K>> for TR<R, K> {
impl<R: Ring> Add<TR<R>> for TR<R> {
type Output = Self;
fn add(self, other: Self) -> Self {
Self(
zip_eq(self.0, other.0)
debug_assert_eq!(self.k, other.k);
Self {
k: self.k,
r: zip_eq(self.r, other.r)
.map(|(s, o)| s + o)
.collect::<Vec<_>>(),
)
}
}
}
impl<R: Ring, const K: usize> Sub<TR<R, K>> for TR<R, K> {
impl<R: Ring> Sub<TR<R>> for TR<R> {
type Output = Self;
fn sub(self, other: Self) -> Self {
Self(zip_eq(self.0, other.0).map(|(s, o)| s - o).collect())
debug_assert_eq!(self.k, other.k);
Self {
k: self.k,
r: zip_eq(self.r, other.r).map(|(s, o)| s - o).collect(),
}
}
}
impl<R: Ring, const K: usize> Neg for TR<R, K> {
impl<R: Ring> Neg for TR<R> {
type Output = Self;
fn neg(self) -> Self::Output {
Self(self.0.iter().map(|&e_i| -e_i).collect())
Self {
k: self.k,
r: self.r.iter().map(|e_i| -e_i.clone()).collect(),
}
}
}
/// for (TR,TR), the Mul operation is defined as the dot product:
/// for A, B \in R^k, result = Σ A_i * B_i \in R
impl<R: Ring, const K: usize> Mul<TR<R, K>> for TR<R, K> {
impl<R: Ring> Mul<TR<R>> for TR<R> {
type Output = R;
fn mul(self, other: Self) -> R {
zip_eq(self.0, other.0).map(|(s, o)| s * o).sum()
debug_assert_eq!(self.k, other.k);
zip_eq(self.r, other.r).map(|(s, o)| s * o).sum()
}
}
impl<R: Ring, const K: usize> Mul<&TR<R, K>> for &TR<R, K> {
impl<R: Ring> Mul<&TR<R>> for &TR<R> {
type Output = R;
fn mul(self, other: &TR<R, K>) -> R {
zip_eq(self.0.clone(), other.0.clone())
fn mul(self, other: &TR<R>) -> R {
debug_assert_eq!(self.k, other.k);
zip_eq(self.r.clone(), other.r.clone())
.map(|(s, o)| s * o)
.sum()
}
@@ -108,15 +135,21 @@ impl<R: Ring, const K: usize> Mul<&TR<R, K>> for &TR<R, K> {
/// for (TR, R), the Mul operation is defined as each element of TR is
/// multiplied by R
impl<R: Ring, const K: usize> 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> Mul<R> for TR<R> {
type Output = TR<R>;
fn mul(self, other: R) -> TR<R> {
Self {
k: self.k,
r: self.r.iter().map(|s| s.clone() * other.clone()).collect(),
}
}
}
impl<R: Ring, const K: usize> 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())
impl<R: Ring> Mul<&R> for &TR<R> {
type Output = TR<R>;
fn mul(self, other: &R) -> TR<R> {
TR::<R> {
k: self.k,
r: self.r.iter().map(|s| s.clone() * other.clone()).collect(),
}
}
}

307
arith/src/zq.rs
View File

@@ -4,41 +4,39 @@ use std::ops::{Add, AddAssign, Div, Mul, Neg, Sub, SubAssign};
/// Z_q, integers modulus q, not necessarily prime
#[derive(Clone, Copy, PartialEq)]
pub struct Zq<const Q: u64>(pub u64);
// WIP
// impl<const Q: u64> From<Vec<u64>> for Vec<Zq<Q>> {
// fn from(v: Vec<u64>) -> Self {
// v.into_iter().map(Zq::new).collect()
// }
// }
pub(crate) fn modulus_u64<const Q: u64>(e: u64) -> u64 {
(e % Q + Q) % Q
pub struct Zq {
pub q: u64,
pub v: u64,
}
impl<const Q: u64> Zq<Q> {
pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>) -> Self {
pub(crate) fn modulus_u64(q: u64, e: u64) -> u64 {
(e % q + q) % q
}
impl Zq {
pub fn rand(mut rng: impl Rng, dist: impl Distribution<f64>, q: u64) -> Self {
// TODO WIP
let r: f64 = dist.sample(&mut rng);
Self::from_f64(r)
// Self::from_u64(r.round() as u64)
Self::from_f64(q, r)
}
pub fn from_u64(e: u64) -> Self {
if e >= Q {
// (e % Q + Q) % Q
return Zq(modulus_u64::<Q>(e));
// return Zq(e % Q);
pub fn from_u64(q: u64, v: u64) -> Self {
if v >= q {
// (v % Q + Q) % Q
return Zq {
q,
v: modulus_u64(q, v),
};
// return Zq(v % Q);
}
Zq(e)
Zq { q, v }
}
pub fn from_f64(e: f64) -> Self {
pub fn from_f64(q: u64, e: f64) -> Self {
// WIP method
let e: i64 = e.round() as i64;
let q = Q as i64;
if e < 0 || e >= q {
return Zq(((e % q + q) % q) as u64);
let q_i64 = q as i64;
if e < 0 || e >= q_i64 {
return Zq::from_u64(q, ((e % q_i64 + q_i64) % q_i64) as u64);
}
Zq(e as u64)
Zq { q, v: e as u64 }
// if e < 0 {
// // dbg!(&e);
@@ -50,15 +48,18 @@ impl<const Q: u64> Zq<Q> {
// }
// Zq(e as u64)
}
pub fn from_bool(b: bool) -> Self {
pub fn from_bool(q: u64, b: bool) -> Self {
if b {
Zq(1)
Zq { q, v: 1 }
} else {
Zq(0)
Zq { q, v: 0 }
}
}
pub fn zero() -> Self {
Self(0u64)
pub fn zero(q: u64) -> Self {
Self { q, v: 0u64 }
}
pub fn one(q: u64) -> Self {
Self { q, v: 1u64 }
}
pub fn square(self) -> Self {
self * self
@@ -66,18 +67,21 @@ impl<const Q: u64> Zq<Q> {
// modular exponentiation
pub fn exp(self, e: Self) -> Self {
// mul-square approach
let mut res = Self(1);
let mut res = Self::one(self.q);
let mut rem = e.clone();
let mut exp = self;
// for rem != Self(0) {
while rem != Self(0) {
while rem != Self::zero(self.q) {
// if odd
// TODO use a more readible expression
if 1 - ((rem.0 & 1) << 1) as i64 == -1 {
// TODO use a more readeable expression
if 1 - ((rem.v & 1) << 1) as i64 == -1 {
res = res * exp;
}
exp = exp.square();
rem = Self(rem.0 >> 1);
rem = Self {
q: self.q,
v: rem.v >> 1,
};
}
res
}
@@ -89,9 +93,9 @@ impl<const Q: u64> Zq<Q> {
// let a = self.0;
// let q = Q;
let mut t = 0;
let mut r = Q;
let mut r = self.q;
let mut new_t = 0;
let mut new_r = self.0.clone();
let mut new_r = self.v.clone();
while new_r != 0 {
let q = r / new_r;
@@ -104,16 +108,16 @@ impl<const Q: u64> Zq<Q> {
// if t < 0 {
// t = t + q;
// }
return Zq::from_u64(t);
return Zq::from_u64(self.q, t);
}
pub fn inv(self) -> Zq<Q> {
let (g, x, _) = Self::egcd(self.0 as i128, Q as i128);
pub fn inv(self) -> Zq {
let (g, x, _) = Self::egcd(self.v as i128, self.q as i128);
if g != 1 {
// None
panic!("E");
} else {
let q = Q as i128;
Zq(((x % q + q) % q) as u64) // TODO maybe just Zq::new(x)
let q = self.q as i128;
Zq::from_u64(self.q, ((x % q + q) % q) as u64) // TODO maybe just Zq::new(x)
}
}
fn egcd(a: i128, b: i128) -> (i128, i128, i128) {
@@ -126,8 +130,11 @@ impl<const Q: u64> Zq<Q> {
}
/// perform the mod switch operation from Q to Q', where Q2=Q'
pub fn mod_switch<const Q2: u64>(&self) -> Zq<Q2> {
Zq::<Q2>::from_u64(((self.0 as f64 * Q2 as f64) / Q as f64).round() as u64)
pub fn mod_switch(&self, q2: u64) -> Zq {
Zq::from_u64(
q2,
((self.v as f64 * q2 as f64) / self.q as f64).round() as u64,
)
}
pub fn decompose(&self, beta: u32, l: u32) -> Vec<Self> {
@@ -138,19 +145,25 @@ impl<const Q: u64> Zq<Q> {
}
}
pub fn decompose_base_beta(&self, beta: u32, l: u32) -> Vec<Self> {
let mut rem: u64 = self.0;
let mut rem: u64 = self.v;
// next if is for cases in which beta does not divide Q (concretely
// beta^l!=Q). round to the nearest multiple of q/beta^l
if rem >= beta.pow(l) as u64 {
// rem = Q - 1 - (Q / beta as u64); // floor
return vec![Zq(beta as u64 - 1); l as usize];
return vec![
Zq {
q: self.q,
v: beta as u64 - 1
};
l as usize
];
}
let mut x: Vec<Self> = vec![];
for i in 1..l + 1 {
let den = Q / beta.pow(i) as u64;
let den = self.q / beta.pow(i) as u64;
let x_i = rem / den; // division between u64 already does floor
x.push(Self::from_u64(x_i));
x.push(Self::from_u64(self.q, x_i));
if x_i != 0 {
rem = rem % den;
}
@@ -161,15 +174,15 @@ impl<const Q: u64> Zq<Q> {
pub fn decompose_base2(&self, l: u32) -> Vec<Self> {
// next if is for cases in which beta does not divide Q (concretely
// beta^l!=Q). round to the nearest multiple of q/beta^l
if self.0 >= 1 << l as u64 {
if self.v >= 1 << l as u64 {
// rem = Q - 1 - (Q / beta as u64); // floor
// (where beta=2)
return vec![Zq(1); l as usize];
return vec![Zq::one(self.q); l as usize];
}
(0..l)
.rev()
.map(|i| Self(((self.0 >> i) & 1) as u64))
.map(|i| Self::from_u64(self.q, ((self.v >> i) & 1) as u64))
.collect()
// naive ver:
@@ -194,114 +207,143 @@ impl<const Q: u64> Zq<Q> {
}
}
impl<const Q: u64> Zq<Q> {
impl Zq {
fn r#mod(self) -> Self {
if self.0 >= Q {
return Zq(self.0 % Q);
if self.v >= self.q {
return Zq::from_u64(self.q, self.v % self.q);
}
self
}
}
impl<const Q: u64> Add<Zq<Q>> for Zq<Q> {
impl Add<Zq> for Zq {
type Output = Self;
fn add(self, rhs: Self) -> Self::Output {
let mut r = self.0 + rhs.0;
if r >= Q {
r -= Q;
}
Zq(r)
}
}
impl<const Q: u64> Add<&Zq<Q>> for &Zq<Q> {
type Output = Zq<Q>;
assert_eq!(self.q, rhs.q);
fn add(self, rhs: &Zq<Q>) -> Self::Output {
let mut r = self.0 + rhs.0;
if r >= Q {
r -= Q;
let mut v = self.v + rhs.v;
if v >= self.q {
v -= self.q;
}
Zq(r)
Zq { q: self.q, v }
}
}
impl<const Q: u64> AddAssign<Zq<Q>> for Zq<Q> {
impl Add<&Zq> for &Zq {
type Output = Zq;
fn add(self, rhs: &Zq) -> Self::Output {
assert_eq!(self.q, rhs.q);
let mut v = self.v + rhs.v;
if v >= self.q {
v -= self.q;
}
Zq { q: self.q, v }
}
}
impl AddAssign<Zq> for Zq {
fn add_assign(&mut self, rhs: Self) {
*self = *self + rhs
}
}
impl<const Q: u64> std::iter::Sum for Zq<Q> {
fn sum<I>(iter: I) -> Self
impl std::iter::Sum for Zq {
fn sum<I>(mut iter: I) -> Self
where
I: Iterator<Item = Self>,
{
iter.fold(Zq(0), |acc, x| acc + x)
let first: Zq = iter.next().unwrap();
iter.fold(first, |acc, x| acc + x)
}
}
impl<const Q: u64> Sub<Zq<Q>> for Zq<Q> {
impl Sub<Zq> for Zq {
type Output = Self;
fn sub(self, rhs: Self) -> Zq<Q> {
if self.0 >= rhs.0 {
Zq(self.0 - rhs.0)
} else {
Zq((Q + self.0) - rhs.0)
}
}
}
impl<const Q: u64> Sub<&Zq<Q>> for &Zq<Q> {
type Output = Zq<Q>;
fn sub(self, rhs: Self) -> Zq {
assert_eq!(self.q, rhs.q);
fn sub(self, rhs: &Zq<Q>) -> Self::Output {
if self.0 >= rhs.0 {
Zq(self.0 - rhs.0)
if self.v >= rhs.v {
Zq {
q: self.q,
v: self.v - rhs.v,
}
} else {
Zq((Q + self.0) - rhs.0)
Zq {
q: self.q,
v: (self.q + self.v) - rhs.v,
}
}
}
}
impl<const Q: u64> SubAssign<Zq<Q>> for Zq<Q> {
impl Sub<&Zq> for &Zq {
type Output = Zq;
fn sub(self, rhs: &Zq) -> Self::Output {
assert_eq!(self.q, rhs.q);
if self.q >= rhs.q {
Zq {
q: self.q,
v: self.v - rhs.v,
}
} else {
Zq {
q: self.q,
v: (self.q + self.v) - rhs.v,
}
}
}
}
impl SubAssign<Zq> for Zq {
fn sub_assign(&mut self, rhs: Self) {
*self = *self - rhs
}
}
impl<const Q: u64> Neg for Zq<Q> {
impl Neg for Zq {
type Output = Self;
fn neg(self) -> Self::Output {
if self.0 == 0 {
if self.v == 0 {
return self;
}
Zq(Q - self.0)
Zq {
q: self.q,
v: self.q - self.v,
}
}
}
impl<const Q: u64> Mul<Zq<Q>> for Zq<Q> {
impl Mul<Zq> for Zq {
type Output = Self;
fn mul(self, rhs: Self) -> Zq<Q> {
fn mul(self, rhs: Self) -> Zq {
assert_eq!(self.q, rhs.q);
// TODO non-naive way
Zq(((self.0 as u128 * rhs.0 as u128) % Q as u128) as u64)
Zq {
q: self.q,
v: ((self.v as u128 * rhs.v as u128) % self.q as u128) as u64,
}
// Zq((self.0 * rhs.0) % Q)
}
}
impl<const Q: u64> Div<Zq<Q>> for Zq<Q> {
impl Div<Zq> for Zq {
type Output = Self;
fn div(self, rhs: Self) -> Zq<Q> {
fn div(self, rhs: Self) -> Zq {
// TODO non-naive way
// Zq((self.0 / rhs.0) % Q)
self * rhs.inv()
}
}
impl<const Q: u64> fmt::Display for Zq<Q> {
impl fmt::Display for Zq {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}", self.0)
write!(f, "{}", self.v)
}
}
impl<const Q: u64> fmt::Debug for Zq<Q> {
impl fmt::Debug for Zq {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{}", self.0)
write!(f, "{}", self.v)
}
}
@@ -312,80 +354,83 @@ mod tests {
#[test]
fn exp() {
const Q: u64 = 1021;
let a = Zq::<Q>(3);
let b = Zq::<Q>(3);
assert_eq!(a.exp(b), Zq(27));
const q: u64 = 1021;
let a = Zq::from_u64(q, 3);
let b = Zq::from_u64(q, 3);
assert_eq!(a.exp(b), Zq::from_u64(q, 27));
let a = Zq::<Q>(1000);
let b = Zq::<Q>(3);
assert_eq!(a.exp(b), Zq(949));
let a = Zq::from_u64(q, 1000);
let b = Zq::from_u64(q, 3);
assert_eq!(a.exp(b), Zq::from_u64(q, 949));
}
#[test]
fn neg() {
const Q: u64 = 1021;
let a = Zq::<Q>::from_f64(101.0);
let b = Zq::<Q>::from_f64(-1.0);
let q: u64 = 1021;
let a = Zq::from_f64(q, 101.0);
let b = Zq::from_f64(q, -1.0);
assert_eq!(-a, a * b);
}
fn recompose<const Q: u64>(beta: u32, l: u32, d: Vec<Zq<Q>>) -> Zq<Q> {
fn recompose(q: u64, beta: u32, l: u32, d: Vec<Zq>) -> Zq {
let mut x = 0u64;
for i in 0..l {
x += d[i as usize].0 * Q / beta.pow(i + 1) as u64;
x += d[i as usize].v * q / beta.pow(i + 1) as u64;
}
Zq::from_u64(x)
Zq::from_u64(q, x)
}
#[test]
fn test_decompose() {
const Q1: u64 = 16;
let q1: u64 = 16;
let beta: u32 = 2;
let l: u32 = 4;
let x = Zq::<Q1>::from_u64(9);
let x = Zq::from_u64(q1, 9);
let d = x.decompose(beta, l);
assert_eq!(recompose::<Q1>(beta, l, d), x);
assert_eq!(recompose(q1, beta, l, d), x);
const Q: u64 = 5u64.pow(3);
let q: u64 = 5u64.pow(3);
let beta: u32 = 5;
let l: u32 = 3;
let dist = Uniform::new(0_u64, Q);
let dist = Uniform::new(0_u64, q);
let mut rng = rand::thread_rng();
for _ in 0..1000 {
let x = Zq::<Q>::from_u64(dist.sample(&mut rng));
let x = Zq::from_u64(q, dist.sample(&mut rng));
let d = x.decompose(beta, l);
assert_eq!(d.len(), l as usize);
assert_eq!(recompose::<Q>(beta, l, d), x);
assert_eq!(recompose(q, beta, l, d), x);
}
}
#[test]
fn test_decompose_approx() {
const Q: u64 = 2u64.pow(4) + 1;
let q: u64 = 2u64.pow(4) + 1;
let beta: u32 = 2;
let l: u32 = 4;
let x = Zq::<Q>::from_u64(16); // in q, but bigger than beta^l
let x = Zq::from_u64(q, 16); // in q, but bigger than beta^l
let d = x.decompose(beta, l);
assert_eq!(d.len(), l as usize);
assert_eq!(recompose::<Q>(beta, l, d), Zq(15));
assert_eq!(recompose(q, beta, l, d), Zq::from_u64(q, 15));
const Q2: u64 = 5u64.pow(3) + 1;
let q2: u64 = 5u64.pow(3) + 1;
let beta: u32 = 5;
let l: u32 = 3;
let x = Zq::<Q2>::from_u64(125); // in q, but bigger than beta^l
let x = Zq::from_u64(q2, 125); // in q, but bigger than beta^l
let d = x.decompose(beta, l);
assert_eq!(d.len(), l as usize);
assert_eq!(recompose::<Q2>(beta, l, d), Zq(124));
assert_eq!(recompose(q2, beta, l, d), Zq::from_u64(q2, 124));
const Q3: u64 = 2u64.pow(16) + 1;
let q3: u64 = 2u64.pow(16) + 1;
let beta: u32 = 2;
let l: u32 = 16;
let x = Zq::<Q3>::from_u64(Q3 - 1); // in q, but bigger than beta^l
let x = Zq::from_u64(q3, q3 - 1); // in q, but bigger than beta^l
let d = x.decompose(beta, l);
assert_eq!(d.len(), l as usize);
assert_eq!(recompose::<Q3>(beta, l, d), Zq(beta.pow(l) as u64 - 1));
assert_eq!(
recompose(q3, beta, l, d),
Zq::from_u64(q3, beta.pow(l) as u64 - 1)
);
}
}