From c65522099b38d40585697a052f7895ffbc30cd47 Mon Sep 17 00:00:00 2001 From: Jean-Philippe Bossuat Date: Mon, 30 Dec 2024 16:20:25 +0100 Subject: [PATCH] Added const for modular reduction, simplfied API --- benches/operations.rs | 42 +++++++- src/dft/ntt.rs | 9 +- src/modulus.rs | 45 +++++--- src/modulus/impl_u64/montgomery.rs | 166 +++++++++++------------------ src/modulus/impl_u64/operations.rs | 60 +++++++---- src/modulus/impl_u64/prime.rs | 33 +++--- src/modulus/montgomery.rs | 11 ++ src/poly.rs | 12 ++- src/ring/impl_u64/ring.rs | 11 +- 9 files changed, 227 insertions(+), 162 deletions(-) diff --git a/benches/operations.rs b/benches/operations.rs index 8b20715..c77ef7e 100644 --- a/benches/operations.rs +++ b/benches/operations.rs @@ -1,6 +1,10 @@ use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use math::ring::Ring; use math::modulus::Operations; +use math::modulus::montgomery::Montgomery; +use math::modulus::{NONE, ONCE}; + +const CHUNK: usize= 8; fn add_vec_unary(c: &mut Criterion) { fn runner(r: Ring) -> Box { @@ -11,9 +15,8 @@ fn add_vec_unary(c: &mut Criterion) { p0.0[i] = i as u64; p1.0[i] = i as u64; } - println!("{}", r.n()); Box::new(move || { - r.modulus.add_vec_unary_assign::<8>(&p0.0, &mut p1.0); + r.modulus.add_vec_unary_assign::(&p0.0, &mut p1.0); }) } @@ -36,5 +39,38 @@ fn add_vec_unary(c: &mut Criterion) { } } -criterion_group!(benches, add_vec_unary); +fn mul_vec_montgomery_external_unary_assign(c: &mut Criterion) { + fn runner(r: Ring) -> Box { + + let mut p0: math::poly::Poly> = r.new_poly_montgomery(); + let mut p1: math::poly::Poly = r.new_poly(); + for i in 0..p0.n(){ + p0.0[i] = r.modulus.montgomery.prepare::(i as u64); + p1.0[i] = i as u64; + } + Box::new(move || { + r.modulus.mul_vec_montgomery_external_unary_assign::(&p0.0, &mut p1.0); + }) + } + + let mut b: criterion::BenchmarkGroup<'_, criterion::measurement::WallTime> = c.benchmark_group("mul_vec_montgomery_external_unary_assign"); + for log_n in 11..17 { + + let n: usize = 1< = Ring::::new(n, q_base, q_power); + let runners = [ + ("prime", { + runner(r) + }), + ]; + for (name, mut runner) in runners { + let id = BenchmarkId::new(name, n); + b.bench_with_input(id, &(), |b, _| b.iter(&mut runner)); + } + } +} + +criterion_group!(benches, add_vec_unary, mul_vec_montgomery_external_unary_assign); criterion_main!(benches); diff --git a/src/dft/ntt.rs b/src/dft/ntt.rs index 21ade41..3b91024 100644 --- a/src/dft/ntt.rs +++ b/src/dft/ntt.rs @@ -3,6 +3,7 @@ use crate::modulus::shoup::Shoup; use crate::modulus::prime::Prime; use crate::modulus::ReduceOnce; use crate::modulus::WordOps; +use crate::modulus::ONCE; use crate::dft::DFT; use itertools::izip; @@ -22,7 +23,7 @@ impl Table< u64> { let psi: u64 = prime.primitive_nth_root(nth_root); - let psi_mont: Montgomery = prime.montgomery.prepare(psi); + let psi_mont: Montgomery = prime.montgomery.prepare::(psi); let psi_inv_mont: Montgomery = prime.montgomery.pow(psi_mont, prime.phi-1); let mut psi_forward_rev: Vec> = vec![Shoup(0, 0); (nth_root >> 1) as usize]; @@ -40,8 +41,8 @@ impl Table< u64> { let i_rev: usize = i.reverse_bits_msb(log_nth_root_half); - prime.montgomery.mul_external_assign(psi_mont, &mut powers_forward); - prime.montgomery.mul_external_assign(psi_inv_mont, &mut powers_backward); + prime.montgomery.mul_external_assign::(psi_mont, &mut powers_forward); + prime.montgomery.mul_external_assign::(psi_inv_mont, &mut powers_backward); psi_forward_rev[i_rev] = prime.shoup.prepare(powers_forward); psi_backward_rev[i_rev] = prime.shoup.prepare(powers_backward); @@ -61,7 +62,7 @@ impl Table< u64> { // Returns n^-1 mod q in Montgomery. fn inv(&self, n:u64) -> Montgomery{ - self.prime.montgomery.pow(self.prime.montgomery.prepare(n), self.prime.phi-1) + self.prime.montgomery.pow(self.prime.montgomery.prepare::(n), self.prime.phi-1) } } diff --git a/src/modulus.rs b/src/modulus.rs index 3fc01c0..3893656 100644 --- a/src/modulus.rs +++ b/src/modulus.rs @@ -4,7 +4,14 @@ pub mod montgomery; pub mod shoup; pub mod impl_u64; +pub type REDUCEMOD = u8; +pub const NONE: REDUCEMOD = 0; +pub const ONCE: REDUCEMOD = 1; +pub const TWICE: REDUCEMOD = 2; +pub const FOURTIMES: REDUCEMOD = 3; +pub const BARRETT: REDUCEMOD = 4; +pub const BARRETTLAZY: REDUCEMOD = 5; pub trait WordOps{ fn log2(self) -> O; @@ -64,47 +71,61 @@ impl ReduceOnce for u64{ #[inline(always)] fn reduce_once_assign(&mut self, q: u64){ debug_assert!(q < 0x8000000000000000, "2q >= 2^64"); - *self = (*self).min(self.wrapping_sub(q)) + *self = *self.min(&mut self.wrapping_sub(q)) } #[inline(always)] fn reduce_once(&self, q:u64) -> u64{ debug_assert!(q < 0x8000000000000000, "2q >= 2^64"); - (*self).min(self.wrapping_sub(q)) + *self.min(&mut self.wrapping_sub(q)) } } pub trait Operations{ // Assigns a + b to c. - fn add_binary_assign(&self, a: &O, b:&O, c: &mut O); + fn add_binary_assign(&self, a: &O, b:&O, c: &mut O); // Assigns a + b to b. - fn add_unary_assign(&self, a: &O, b: &mut O); + fn add_unary_assign(&self, a: &O, b: &mut O); // Assigns a[i] + b[i] to c[i] - fn add_vec_binary_assign(&self, a: &[O], b:&[O], c: &mut [O]); + fn add_vec_binary_assign(&self, a: &[O], b:&[O], c: &mut [O]); // Assigns a[i] + b[i] to b[i] - fn add_vec_unary_assign(&self, a: &[O], b: &mut [O]); + fn add_vec_unary_assign(&self, a: &[O], b: &mut [O]); // Assigns a - b to c. - fn sub_binary_assign(&self, a: &O, b:&O, c: &mut O); + fn sub_binary_assign(&self, a: &O, b:&O, c: &mut O); // Assigns b - a to b. - fn sub_unary_assign(&self, a: &O, b: &mut O); + fn sub_unary_assign(&self, a: &O, b: &mut O); // Assigns a[i] - b[i] to c[i] - fn sub_vec_binary_assign(&self, a: &[O], b:&[O], c: &mut [O]); + fn sub_vec_binary_assign(&self, a: &[O], b:&[O], c: &mut [O]); // Assigns a[i] - b[i] to b[i] - fn sub_vec_unary_assign(&self, a: &[O], b: &mut [O]); + fn sub_vec_unary_assign(&self, a: &[O], b: &mut [O]); // Assigns -a to a. - fn neg_assign(&self, a:&mut O); + fn neg_assign(&self, a:&mut O); // Assigns -a[i] to a[i]. - fn neg_vec_assign(&self, a: &mut [O]); + fn neg_vec_assign(&self, a: &mut [O]); + + // Assigns a * b to c. + fn mul_montgomery_external_binary_assign(&self, a:&montgomery::Montgomery, b:&O, c: &mut O); + + // Assigns a * b to b. + fn mul_montgomery_external_unary_assign(&self, a:&montgomery::Montgomery, b:&mut O); + + // Assigns a[i] * b[i] to c[i]. + fn mul_vec_montgomery_external_binary_assign(&self, a:&[montgomery::Montgomery], b:&[O], c: &mut [O]); + + // Assigns a[i] * b[i] to b[i]. + fn mul_vec_montgomery_external_unary_assign(&self, a:&[montgomery::Montgomery], b:&mut [O]); } + + diff --git a/src/modulus/impl_u64/montgomery.rs b/src/modulus/impl_u64/montgomery.rs index 6119d25..3522f32 100644 --- a/src/modulus/impl_u64/montgomery.rs +++ b/src/modulus/impl_u64/montgomery.rs @@ -1,7 +1,8 @@ use crate::modulus::ReduceOnce; use crate::modulus::montgomery::{MontgomeryPrecomp, Montgomery}; -use crate::modulus::barrett::{BarrettPrecomp}; +use crate::modulus::barrett::BarrettPrecomp; +use crate::modulus::{REDUCEMOD, NONE, ONCE, TWICE, FOURTIMES, BARRETT, BARRETTLAZY}; extern crate test; /// MontgomeryPrecomp is a set of methods implemented for MontgomeryPrecomp @@ -20,14 +21,16 @@ impl MontgomeryPrecomp{ q_pow = q_pow.wrapping_mul(q_pow); } let mut precomp = Self{ - q: q, + q: q, + two_q: q<<1, + four_q: q<<2, barrett: BarrettPrecomp::new(q), q_inv: q_inv, one: Montgomery(0), minus_one: Montgomery(0), }; - precomp.one = precomp.prepare(1); + precomp.one = precomp.prepare::(1); precomp.minus_one = Montgomery(q-precomp.one.value()); precomp @@ -45,120 +48,91 @@ impl MontgomeryPrecomp{ self.minus_one } + /// Applies a modular reduction on x based on REDUCE: + /// - LAZY: no modular reduction. + /// - ONCE: subtracts q if x >= q. + /// - FULL: maps x to x mod q using Barrett reduction. + #[inline(always)] + pub fn reduce(&self, x: u64) -> u64{ + let mut r: u64 = x; + self.reduce_assign::(&mut r); + r + } + + /// Applies a modular reduction on x based on REDUCE: + /// - LAZY: no modular reduction. + /// - ONCE: subtracts q if x >= q. + /// - FULL: maps x to x mod q using Barrett reduction. + #[inline(always)] + pub fn reduce_assign(&self, x: &mut u64){ + match REDUCE { + NONE =>{}, + ONCE =>{x.reduce_once_assign(self.q)}, + TWICE=>{x.reduce_once_assign(self.two_q)}, + FOURTIMES =>{x.reduce_once_assign(self.four_q)}, + BARRETT =>{self.barrett.reduce_assign(x)}, + BARRETTLAZY =>{self.barrett.reduce_assign(x)}, + _ => unreachable!("invalid REDUCE argument") + } + } + /// Returns lhs * 2^64 mod q as a Montgomery. #[inline(always)] - pub fn prepare(&self, lhs: u64) -> Montgomery{ + pub fn prepare(&self, lhs: u64) -> Montgomery{ let mut rhs = Montgomery(0); - self.prepare_assign(lhs, &mut rhs); + self.prepare_assign::(lhs, &mut rhs); rhs } /// Assigns lhs * 2^64 mod q to rhs. #[inline(always)] - pub fn prepare_assign(&self, lhs: u64, rhs: &mut Montgomery){ - self.prepare_lazy_assign(lhs, rhs); - rhs.value_mut().reduce_once_assign(self.q); - } - - /// Returns lhs * 2^64 mod q in range [0, 2q-1] as a Montgomery. - #[inline(always)] - pub fn prepare_lazy(&self, lhs: u64) -> Montgomery{ - let mut rhs = Montgomery(0); - self.prepare_lazy_assign(lhs, &mut rhs); - rhs - } - - /// Assigns lhs * 2^64 mod q in range [0, 2q-1] to rhs. - #[inline(always)] - pub fn prepare_lazy_assign(&self, lhs: u64, rhs: &mut Montgomery){ + pub fn prepare_assign(&self, lhs: u64, rhs: &mut Montgomery){ let (_, mhi) = lhs.widening_mul(*self.barrett.value_lo()); *rhs = Montgomery((lhs.wrapping_mul(*self.barrett.value_hi()).wrapping_add(mhi)).wrapping_mul(self.q).wrapping_neg()); + self.reduce_assign::(rhs.value_mut()); } /// Returns lhs * (2^64)^-1 mod q as a u64. #[inline(always)] - pub fn unprepare(&self, lhs: Montgomery) -> u64{ + pub fn unprepare(&self, lhs: Montgomery) -> u64{ let mut rhs = 0u64; - self.unprepare_assign(lhs, &mut rhs); + self.unprepare_assign::(lhs, &mut rhs); rhs } /// Assigns lhs * (2^64)^-1 mod q to rhs. #[inline(always)] - pub fn unprepare_assign(&self, lhs: Montgomery, rhs: &mut u64){ - self.unprepare_lazy_assign(lhs, rhs); - rhs.reduce_once_assign(self.q); - - } - - /// Returns lhs * (2^64)^-1 mod q in range [0, 2q-1]. - #[inline(always)] - pub fn unprepare_lazy(&self, lhs: Montgomery) -> u64{ - let mut rhs = 0u64; - self.unprepare_lazy_assign(lhs, &mut rhs); - rhs - } - - /// Assigns lhs * (2^64)^-1 mod q in range [0, 2q-1] to rhs. - #[inline(always)] - pub fn unprepare_lazy_assign(&self, lhs: Montgomery, rhs: &mut u64){ + pub fn unprepare_assign(&self, lhs: Montgomery, rhs: &mut u64){ let (_, r) = self.q.widening_mul(lhs.value().wrapping_mul(self.q_inv)); - *rhs = self.q - r + *rhs = self.reduce::(self.q.wrapping_sub(r)); } /// Returns lhs * rhs * (2^{64})^-1 mod q. #[inline(always)] - pub fn mul_external(&self, lhs: Montgomery, rhs: u64) -> u64{ - let mut r = self.mul_external_lazy(lhs, rhs); - r.reduce_once_assign(self.q); + pub fn mul_external(&self, lhs: Montgomery, rhs: u64) -> u64{ + let mut r: u64 = rhs; + self.mul_external_assign::(lhs, &mut r); r } /// Assigns lhs * rhs * (2^{64})^-1 mod q to rhs. #[inline(always)] - pub fn mul_external_assign(&self, lhs: Montgomery, rhs: &mut u64){ - self.mul_external_lazy_assign(lhs, rhs); - rhs.reduce_once_assign(self.q); - } - - /// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1]. - #[inline(always)] - pub fn mul_external_lazy(&self, lhs: Montgomery, rhs: u64) -> u64{ - let mut result: u64 = rhs; - self.mul_external_lazy_assign(lhs, &mut result); - result - } - - /// Assigns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1] to rhs. - #[inline(always)] - pub fn mul_external_lazy_assign(&self, lhs: Montgomery, rhs: &mut u64){ + pub fn mul_external_assign(&self, lhs: Montgomery, rhs: &mut u64){ let (mlo, mhi) = lhs.value().widening_mul(*rhs); let (_, hhi) = self.q.widening_mul(mlo.wrapping_mul(self.q_inv)); - *rhs = mhi.wrapping_add(self.q - hhi) + *rhs = self.reduce::(mhi.wrapping_sub(hhi).wrapping_add(self.q)); } /// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1]. #[inline(always)] - pub fn mul_internal(&self, lhs: Montgomery, rhs: Montgomery) -> Montgomery{ - Montgomery(self.mul_external(lhs, *rhs.value())) + pub fn mul_internal(&self, lhs: Montgomery, rhs: Montgomery) -> Montgomery{ + Montgomery(self.mul_external::(lhs, *rhs.value())) } /// Assigns lhs * rhs * (2^{64})^-1 mod q to rhs. #[inline(always)] - pub fn mul_internal_assign(&self, lhs: Montgomery, rhs: &mut Montgomery){ - self.mul_external_assign(lhs, rhs.value_mut()); - } - - /// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1]. - #[inline(always)] - pub fn mul_internal_lazy(&self, lhs: Montgomery, rhs: Montgomery) -> Montgomery{ - Montgomery(self.mul_external_lazy(lhs, *rhs.value())) - } - - /// Assigns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1] to rhs. - #[inline(always)] - pub fn mul_internal_lazy_assign(&self, lhs: Montgomery, rhs: &mut Montgomery){ - self.mul_external_lazy_assign(lhs, rhs.value_mut()); + pub fn mul_internal_assign(&self, lhs: Montgomery, rhs: &mut Montgomery){ + self.mul_external_assign::(lhs, rhs.value_mut()); } #[inline(always)] @@ -174,27 +148,11 @@ impl MontgomeryPrecomp{ /// Assigns lhs + rhs - q if (lhs + rhs) >= q to rhs. #[inline(always)] - pub fn add_internal_reduce_once_assign(&self, lhs: Montgomery, rhs: &mut Montgomery){ + pub fn add_internal_reduce_once_assign(&self, lhs: Montgomery, rhs: &mut Montgomery){ self.add_internal_lazy_assign(lhs, rhs); rhs.value_mut().reduce_once_assign(self.q); } - #[inline(always)] - pub fn reduce(&self, lhs: u64) -> u64{ - self.barrett.reduce(lhs) - } - - /// Returns lhs mod q in range [0, 2q-1]. - #[inline(always)] - pub fn reduce_lazy(&self, lhs: u64) -> u64{ - self.barrett.reduce_lazy(lhs) - } - - #[inline(always)] - pub fn reduce_assign(&self, lhs: &mut u64){ - self.barrett.reduce_assign(lhs) - } - /// Returns lhs mod q in range [0, 2q-1]. #[inline(always)] pub fn reduce_lazy_assign(&self, lhs: &mut u64){ @@ -209,9 +167,9 @@ impl MontgomeryPrecomp{ let mut i: u64 = exponent; while i > 0{ if i & 1 == 1{ - self.mul_internal_assign(x_mut, &mut y); + self.mul_internal_assign::(x_mut, &mut y); } - self.mul_internal_assign(x_mut, &mut x_mut); + self.mul_internal_assign::(x_mut, &mut x_mut); i >>= 1; } @@ -227,16 +185,16 @@ impl MontgomeryPrecomp{ fn pow(x:u64, exponent:u64, q:u64) -> u64{ let montgomery: MontgomeryPrecomp = MontgomeryPrecomp::::new(q); let mut y_mont: Montgomery = montgomery.one(); - let mut x_mont: Montgomery = montgomery.prepare(x); + let mut x_mont: Montgomery = montgomery.prepare::(x); while exponent > 0{ if exponent & 1 == 1{ - montgomery.mul_internal_assign(x_mont, &mut y_mont); + montgomery.mul_internal_assign::(x_mont, &mut y_mont); } - montgomery.mul_internal_assign(x_mont, &mut x_mont); + montgomery.mul_internal_assign::(x_mont, &mut x_mont); } - montgomery.unprepare(y_mont) + montgomery.unprepare::(y_mont) } #[cfg(test)] @@ -251,8 +209,8 @@ mod tests { let m_precomp = montgomery::MontgomeryPrecomp::new(q); let x: u64 = 0x5f876e514845cc8b; let y: u64 = 0xad726f98f24a761a; - let y_mont = m_precomp.prepare(y); - assert!(m_precomp.mul_external(y_mont, x) == (x as u128 * y as u128 % q as u128) as u64); + let y_mont = m_precomp.prepare::(y); + assert!(m_precomp.mul_external::(y_mont, x) == (x as u128 * y as u128 % q as u128) as u64); } #[bench] @@ -261,7 +219,7 @@ mod tests { let m_precomp = montgomery::MontgomeryPrecomp::new(q); let mut x: u64 = 0x5f876e514845cc8b; let y: u64 = 0xad726f98f24a761a; - let y_mont = m_precomp.prepare(y); - b.iter(|| m_precomp.mul_external_assign(y_mont, &mut x)); + let y_mont = m_precomp.prepare::(y); + b.iter(|| m_precomp.mul_external_assign::(y_mont, &mut x)); } } \ No newline at end of file diff --git a/src/modulus/impl_u64/operations.rs b/src/modulus/impl_u64/operations.rs index 05215cd..9d2c4e0 100644 --- a/src/modulus/impl_u64/operations.rs +++ b/src/modulus/impl_u64/operations.rs @@ -2,59 +2,83 @@ use crate::modulus::Operations; use crate::modulus::prime::Prime; use crate::modulus::ReduceOnce; +use crate::modulus::montgomery::Montgomery; +use crate::modulus::{REDUCEMOD, NONE, ONCE, BARRETT, BARRETTLAZY}; use crate::{apply_unary, apply_binary, apply_ternary}; use itertools::izip; impl Operations for Prime{ #[inline(always)] - fn add_binary_assign(&self, a: &u64, b: &u64, c: &mut u64){ - *c = a.wrapping_add(*b).reduce_once(self.q); + fn add_binary_assign(&self, a: &u64, b: &u64, c: &mut u64){ + *c = a.wrapping_add(*b); + self.montgomery.reduce_assign::(c); } #[inline(always)] - fn add_unary_assign(&self, a: &u64, b: &mut u64){ - *b = a.wrapping_add(*b).reduce_once(self.q); + fn add_unary_assign(&self, a: &u64, b: &mut u64){ + *b = a.wrapping_add(*b); + self.montgomery.reduce_assign::(b); } #[inline(always)] - fn add_vec_binary_assign(&self, a: &[u64], b:&[u64], c:&mut [u64]){ - apply_ternary!(self, Self::add_binary_assign, a, b, c, CHUNK); + fn add_vec_binary_assign(&self, a: &[u64], b:&[u64], c:&mut [u64]){ + apply_ternary!(self, Self::add_binary_assign::, a, b, c, CHUNK); } #[inline(always)] - fn add_vec_unary_assign(&self, a: &[u64], b:&mut [u64]){ - apply_binary!(self, Self::add_unary_assign, a, b, CHUNK); + fn add_vec_unary_assign(&self, a: &[u64], b:&mut [u64]){ + apply_binary!(self, Self::add_unary_assign::, a, b, CHUNK); } #[inline(always)] - fn sub_binary_assign(&self, a: &u64, b: &u64, c: &mut u64){ + fn sub_binary_assign(&self, a: &u64, b: &u64, c: &mut u64){ *c = a.wrapping_add(self.q.wrapping_sub(*b)).reduce_once(self.q); } #[inline(always)] - fn sub_unary_assign(&self, a: &u64, b: &mut u64){ + fn sub_unary_assign(&self, a: &u64, b: &mut u64){ *b = a.wrapping_add(self.q.wrapping_sub(*b)).reduce_once(self.q); } #[inline(always)] - fn sub_vec_binary_assign(&self, a: &[u64], b:&[u64], c:&mut [u64]){ - apply_ternary!(self, Self::sub_binary_assign, a, b, c, CHUNK); + fn sub_vec_binary_assign(&self, a: &[u64], b:&[u64], c:&mut [u64]){ + apply_ternary!(self, Self::sub_binary_assign::, a, b, c, CHUNK); } #[inline(always)] - fn sub_vec_unary_assign(&self, a: &[u64], b:&mut [u64]){ - apply_binary!(self, Self::sub_unary_assign, a, b, CHUNK); + fn sub_vec_unary_assign(&self, a: &[u64], b:&mut [u64]){ + apply_binary!(self, Self::sub_unary_assign::, a, b, CHUNK); } #[inline(always)] - fn neg_assign(&self, a: &mut u64){ + fn neg_assign(&self, a: &mut u64){ *a = self.q.wrapping_sub(*a); + self.montgomery.reduce_assign::(a) } #[inline(always)] - fn neg_vec_assign(&self, a: &mut [u64]){ - apply_unary!(self, Self::neg_assign, a, CHUNK); + fn neg_vec_assign(&self, a: &mut [u64]){ + apply_unary!(self, Self::neg_assign::, a, CHUNK); + } + + #[inline(always)] + fn mul_montgomery_external_binary_assign(&self, a:& Montgomery, b:&u64, c: &mut u64){ + *c = self.montgomery.mul_external::(*a, *b); + } + + #[inline(always)] + fn mul_montgomery_external_unary_assign(&self, lhs:&Montgomery, rhs:&mut u64){ + *rhs = self.montgomery.mul_external::(*lhs, *rhs); + } + + #[inline(always)] + fn mul_vec_montgomery_external_binary_assign(&self, a:& [Montgomery], b:&[u64], c: &mut [u64]){ + apply_ternary!(self, Self::mul_montgomery_external_binary_assign::, a, b, c, CHUNK); + } + + #[inline(always)] + fn mul_vec_montgomery_external_unary_assign(&self, a:&[Montgomery], b:&mut [u64]){ + apply_binary!(self, Self::mul_montgomery_external_unary_assign::, a, b, CHUNK); } } - diff --git a/src/modulus/impl_u64/prime.rs b/src/modulus/impl_u64/prime.rs index 9e17550..c7100f4 100644 --- a/src/modulus/impl_u64/prime.rs +++ b/src/modulus/impl_u64/prime.rs @@ -1,6 +1,7 @@ use crate::modulus::prime::Prime; use crate::modulus::montgomery::{Montgomery, MontgomeryPrecomp}; use crate::modulus::shoup::{ShoupPrecomp}; +use crate::modulus::ONCE; use primality_test::is_prime; use prime_factorization::Factorization; @@ -63,20 +64,20 @@ impl Prime{ #[inline(always)] pub fn pow(&self, x: u64, exponent: u64) -> u64{ let mut y_mont: Montgomery = self.montgomery.one(); - let mut x_mont: Montgomery = self.montgomery.prepare(x); + let mut x_mont: Montgomery = self.montgomery.prepare::(x); let mut i: u64 = exponent; while i > 0{ if i & 1 == 1{ - self.montgomery.mul_internal_assign(x_mont, &mut y_mont); + self.montgomery.mul_internal_assign::(x_mont, &mut y_mont); } - self.montgomery.mul_internal_assign(x_mont, &mut x_mont); + self.montgomery.mul_internal_assign::(x_mont, &mut x_mont); i >>= 1; } - self.montgomery.unprepare(y_mont) + self.montgomery.unprepare::(y_mont) } /// Returns x^-1 mod q. @@ -164,9 +165,7 @@ impl Prime{ if q_base != 1{ panic!("invalid factor list: does not fully divide q_base: q_base % (all factors) = {}", q_base) } - } - - + } } /// Returns (psi + a * q_base)^{nth_root} = 1 mod q = q_base^q_power given psi^{nth_root} = 1 mod q_base. @@ -174,23 +173,23 @@ impl Prime{ fn hensel_lift(&self, psi: u64, nth_root: u64) -> u64{ assert!(Pow(psi, nth_root, self.q_base)==1, "invalid argument psi: psi^nth_root = {} != 1", Pow(psi, nth_root, self.q_base)); - let mut psi_mont: Montgomery = self.montgomery.prepare(psi); - let nth_root_mont: Montgomery = self.montgomery.prepare(nth_root); + let mut psi_mont: Montgomery = self.montgomery.prepare::(psi); + let nth_root_mont: Montgomery = self.montgomery.prepare::(nth_root); for _i in 1..self.q_power{ let psi_pow: Montgomery = self.montgomery.pow(psi_mont, nth_root-1); - let num: Montgomery = Montgomery(self.montgomery.one().value() + self.q - self.montgomery.mul_internal(psi_pow, psi_mont).value()); + let num: Montgomery = Montgomery(self.montgomery.one().value() + self.q - self.montgomery.mul_internal::(psi_pow, psi_mont).value()); - let mut den: Montgomery = self.montgomery.mul_internal(nth_root_mont, psi_pow); + let mut den: Montgomery = self.montgomery.mul_internal::(nth_root_mont, psi_pow); den = self.montgomery.pow(den, self.phi-1); - psi_mont = self.montgomery.add_internal(psi_mont, self.montgomery.mul_internal(num, den)); + psi_mont = self.montgomery.add_internal(psi_mont, self.montgomery.mul_internal::(num, den)); } - self.montgomery.unprepare(psi_mont) + self.montgomery.unprepare::(psi_mont) } } @@ -201,17 +200,17 @@ impl Prime{ pub fn Pow(x:u64, exponent:u64, q:u64) -> u64{ let montgomery: MontgomeryPrecomp = MontgomeryPrecomp::::new(q); let mut y_mont: Montgomery = montgomery.one(); - let mut x_mont: Montgomery = montgomery.prepare(x); + let mut x_mont: Montgomery = montgomery.prepare::(x); let mut i: u64 = exponent; while i > 0{ if i & 1 == 1{ - montgomery.mul_internal_assign(x_mont, &mut y_mont); + montgomery.mul_internal_assign::(x_mont, &mut y_mont); } - montgomery.mul_internal_assign(x_mont, &mut x_mont); + montgomery.mul_internal_assign::(x_mont, &mut x_mont); i >>= 1; } - montgomery.unprepare(y_mont) + montgomery.unprepare::(y_mont) } \ No newline at end of file diff --git a/src/modulus/montgomery.rs b/src/modulus/montgomery.rs index 36d4e3a..9435015 100644 --- a/src/modulus/montgomery.rs +++ b/src/modulus/montgomery.rs @@ -23,11 +23,22 @@ impl Montgomery{ } } +/// Default instantiation. +impl Default for Montgomery where O:Default { + fn default() -> Self { + Self { + 0: O::default(), + } + } +} + /// MontgomeryPrecomp is a generic struct storing /// precomputations for Montgomery arithmetic. #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct MontgomeryPrecomp{ pub q: O, + pub two_q: O, + pub four_q: O, pub barrett: BarrettPrecomp, pub q_inv: O, pub one: Montgomery, diff --git a/src/poly.rs b/src/poly.rs index 1cf8541..c7cccd3 100644 --- a/src/poly.rs +++ b/src/poly.rs @@ -3,16 +3,22 @@ pub mod poly; #[derive(Clone, Debug, PartialEq, Eq)] pub struct Poly(pub Vec); -impl Poly{ +impl Polywhere + O: Default + Clone, + { pub fn new(n: usize) -> Self{ - Self(vec![0u64;n]) + Self(vec![O::default();n]) + } + + pub fn new_montgomery(n: usize) -> Self{ + Self(vec![O::default();n]) } pub fn buffer_size(&self) -> usize{ return self.0.len() } - pub fn from_buffer(&mut self, n: usize, buf: &mut [u64]){ + pub fn from_buffer(&mut self, n: usize, buf: &mut [O]){ assert!(buf.len() >= n, "invalid buffer: buf.len()={} < n={}", buf.len(), n); self.0 = Vec::from(&buf[..n]); } diff --git a/src/ring/impl_u64/ring.rs b/src/ring/impl_u64/ring.rs index debdcf3..61d566b 100644 --- a/src/ring/impl_u64/ring.rs +++ b/src/ring/impl_u64/ring.rs @@ -1,6 +1,7 @@ use crate::ring::Ring; use crate::dft::ntt::Table; use crate::modulus::prime::Prime; +use crate::modulus::montgomery::Montgomery; use crate::poly::Poly; impl Ring{ @@ -17,7 +18,15 @@ impl Ring{ return self.n } + fn new_poly_core(&self) -> Poly where O: Default + Clone { + Poly::::new(self.n()) + } + pub fn new_poly(&self) -> Poly{ - Poly::::new(self.n()) + self.new_poly_core::() + } + + pub fn new_poly_montgomery(&self) -> Poly>{ + self.new_poly_core::>() } } \ No newline at end of file