From ee96c2f9046d95a609ff256551a1529c486a5a14 Mon Sep 17 00:00:00 2001 From: Jean-Philippe Bossuat Date: Wed, 4 Dec 2024 12:53:13 +0100 Subject: [PATCH] Added base for Montgomery arithmetic --- .gitignore | 1 + Cargo.lock | 16 ++++ Cargo.toml | 7 ++ src/lib.rs | 3 + src/modulus.rs | 24 +++++ src/modulus/barrett.rs | 19 ++++ src/modulus/montgomery.rs | 113 ++++++++++++++++++++++++ src/modulus/prime.rs | 19 ++++ src/modulus/prime_generation.rs | 151 ++++++++++++++++++++++++++++++++ 9 files changed, 353 insertions(+) create mode 100644 .gitignore create mode 100644 Cargo.lock create mode 100644 Cargo.toml create mode 100644 src/lib.rs create mode 100644 src/modulus.rs create mode 100644 src/modulus/barrett.rs create mode 100644 src/modulus/montgomery.rs create mode 100644 src/modulus/prime.rs create mode 100644 src/modulus/prime_generation.rs diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..ea8c4bf --- /dev/null +++ b/.gitignore @@ -0,0 +1 @@ +/target diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000..ca364c3 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,16 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "math" +version = "0.1.0" +dependencies = [ + "primality-test", +] + +[[package]] +name = "primality-test" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98439e9658b9548a33abdab8c82532554dc08e49ddc5398a9262222fb360ae24" diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000..ac36730 --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,7 @@ +[package] +name = "math" +version = "0.1.0" +edition = "2021" + +[dependencies] +primality-test = "0.3.0" diff --git a/src/lib.rs b/src/lib.rs new file mode 100644 index 0000000..7ab95ca --- /dev/null +++ b/src/lib.rs @@ -0,0 +1,3 @@ +#![feature(bigint_helper_methods)] + +pub mod modulus; \ No newline at end of file diff --git a/src/modulus.rs b/src/modulus.rs new file mode 100644 index 0000000..581f1d4 --- /dev/null +++ b/src/modulus.rs @@ -0,0 +1,24 @@ +pub(crate) mod prime; +pub(crate) mod montgomery; +pub(crate) mod barrett; + +trait ReduceOnce{ + fn reduce_once_assign(&mut self, q: O); + fn reduce_once(&self, q:O) -> O; +} + +impl ReduceOnce for u64{ + fn reduce_once_assign(&mut self, q: u64){ + if *self >= q{ + *self -= q + } + } + + fn reduce_once(&self, q:u64) -> u64{ + if *self >= q { + *self - q + } else { + *self + } + } +} \ No newline at end of file diff --git a/src/modulus/barrett.rs b/src/modulus/barrett.rs new file mode 100644 index 0000000..afa7283 --- /dev/null +++ b/src/modulus/barrett.rs @@ -0,0 +1,19 @@ +pub struct BarrettPrecomp(O, O); + +impl BarrettPrecomp{ + + #[inline(always)] + pub fn new(a:O, b: O) -> Self{ + Self(a, b) + } + + #[inline(always)] + pub fn value_hi(&self) -> &O{ + &self.0 + } + + #[inline(always)] + pub fn value_lo(&self) -> &O{ + &self.1 + } +} \ No newline at end of file diff --git a/src/modulus/montgomery.rs b/src/modulus/montgomery.rs new file mode 100644 index 0000000..9e976ea --- /dev/null +++ b/src/modulus/montgomery.rs @@ -0,0 +1,113 @@ +use crate::modulus::barrett::BarrettPrecomp; +use crate::modulus::ReduceOnce; + +pub struct Montgomery(O); + + +impl Montgomery{ + + #[inline(always)] + pub fn new(lhs: O) -> Self{ + Self(lhs) + } + + #[inline(always)] + pub fn value(&self) -> &O{ + &self.0 + } + + pub fn value_mut(&mut self) -> &mut O{ + &mut self.0 + } +} + +pub struct MontgomeryPrecomp{ + q: O, + q_barrett: BarrettPrecomp, + q_inv: O, +} + +impl MontgomeryPrecomp{ + + #[inline(always)] + fn new(&self, q: u64) -> MontgomeryPrecomp{ + let mut r: u64 = 1; + let mut q_pow = q; + for _i in 0..63{ + r = r.wrapping_mul(r); + q_pow = q_pow.wrapping_mul(q_pow) + } + Self{ q: q, q_barrett: BarrettPrecomp::new(q, q), q_inv: q_pow} + } + + #[inline(always)] + fn prepare(&self, lhs: u64) -> Montgomery{ + let mut rhs = Montgomery(0); + self.prepare_assign(lhs, &mut rhs); + rhs + } + + fn prepare_assign(&self, lhs: u64, rhs: &mut Montgomery){ + self.prepare_lazy_assign(lhs, rhs); + rhs.value_mut().reduce_once_assign(self.q); + } + + #[inline(always)] + fn prepare_lazy(&self, lhs: u64) -> Montgomery{ + let mut rhs = Montgomery(0); + self.prepare_lazy_assign(lhs, &mut rhs); + rhs + } + + fn prepare_lazy_assign(&self, lhs: u64, rhs: &mut Montgomery){ + let (mhi, _) = lhs.widening_mul(*self.q_barrett.value_lo()); + *rhs = Montgomery((lhs.wrapping_mul(*self.q_barrett.value_hi()).wrapping_add(mhi)).wrapping_mul(self.q).wrapping_neg()); + } + + #[inline(always)] + fn mul_external(&self, lhs: Montgomery, rhs: u64) -> u64{ + let mut r = self.mul_external_lazy(lhs, rhs); + r.reduce_once_assign(self.q); + r + } + + #[inline(always)] + fn mul_external_assign(&self, lhs: Montgomery, rhs: &mut u64){ + self.mul_external_lazy_assign(lhs, rhs); + rhs.reduce_once_assign(self.q); + } + + #[inline(always)] + fn mul_external_lazy(&self, lhs: Montgomery, rhs: u64) -> u64{ + let mut result = rhs; + self.mul_external_lazy_assign(lhs, &mut result); + result + } + + #[inline(always)] + fn mul_external_lazy_assign(&self, lhs: Montgomery, rhs: &mut u64){ + let (mhi, mlo) = lhs.value().widening_mul(*rhs); + let (hhi, _) = self.q.widening_mul(mlo * self.q_inv); + *rhs = mhi - hhi + self.q + } + + #[inline(always)] + fn mul_internal(&self, lhs: Montgomery, rhs: Montgomery) -> Montgomery{ + Montgomery(self.mul_external(lhs, *rhs.value())) + } + + #[inline(always)] + fn mul_internal_assign(&self, lhs: Montgomery, rhs: &mut Montgomery){ + self.mul_external_assign(lhs, rhs.value_mut()); + } + + #[inline(always)] + fn mul_internal_lazy(&self, lhs: Montgomery, rhs: Montgomery) -> Montgomery{ + Montgomery(self.mul_external_lazy(lhs, *rhs.value())) + } + + #[inline(always)] + fn mul_internal_lazy_assign(&self, lhs: Montgomery, rhs: &mut Montgomery){ + self.mul_external_lazy_assign(lhs, rhs.value_mut()); + } +} \ No newline at end of file diff --git a/src/modulus/prime.rs b/src/modulus/prime.rs new file mode 100644 index 0000000..70a3f53 --- /dev/null +++ b/src/modulus/prime.rs @@ -0,0 +1,19 @@ +use primality_test::is_prime; + +pub struct Prime { + q: u64, +} + +impl Prime { + pub fn new(q: u64) -> Self{ + assert!(is_prime(q) && q > 2); + Self::new_unchecked(q) + } + + pub fn new_unchecked(q: u64) -> Self { + assert!(q.next_power_of_two().ilog2() <= 61); + Self { + q, + } + } +} \ No newline at end of file diff --git a/src/modulus/prime_generation.rs b/src/modulus/prime_generation.rs new file mode 100644 index 0000000..0036554 --- /dev/null +++ b/src/modulus/prime_generation.rs @@ -0,0 +1,151 @@ +use crate::modulus::prime; + +use prime::Prime; +use primality_test::is_prime; + +pub struct NTTFriendlyPrimesGenerator{ + size: f64, + next_prime: u64, + prev_prime: u64, + nth_root: u64, + check_next_prime: bool, + check_prev_prime: bool, +} + +impl NTTFriendlyPrimesGenerator { + pub fn new(bit_size: u64, nth_root: u64) -> Self{ + let mut check_next_prime: bool = true; + let mut check_prev_prime: bool = true; + let next_prime = (1< 0xffff_ffff_ffff_ffff-nth_root{ + check_next_prime = false; + } + + if prev_prime < nth_root{ + check_prev_prime = false + } + + prev_prime -= nth_root; + + Self{ + size: bit_size as f64, + check_next_prime, + check_prev_prime, + nth_root, + next_prime, + prev_prime, + } + } + + pub fn next_upstream_primes(&mut self, k: usize) -> Vec{ + let mut primes: Vec = Vec::with_capacity(k); + for i in 0..k{ + primes.push(self.next_upstream_prime()) + } + primes + } + + pub fn next_downstream_primes(&mut self, k: usize) -> Vec{ + let mut primes: Vec = Vec::with_capacity(k); + for i in 0..k{ + primes.push(self.next_downstream_prime()) + } + primes + } + + pub fn next_alternating_primes(&mut self, k: usize) -> Vec{ + let mut primes: Vec = Vec::with_capacity(k); + for i in 0..k{ + primes.push(self.next_alternating_prime()) + } + primes + } + + pub fn next_upstream_prime(&mut self) -> Prime{ + loop { + if self.check_next_prime{ + if (self.next_prime as f64).log2() - self.size >= 0.5 || self.next_prime > 0xffff_ffff_ffff_ffff-self.nth_root{ + self.check_next_prime = false; + panic!("prime list for upstream primes is exhausted (overlap with next bit-size or prime > 2^64)"); + } + }else{ + if is_prime(self.next_prime) { + let prime = Prime::new_unchecked(self.next_prime); + self.next_prime += self.nth_root; + return prime + } + self.next_prime += self.nth_root; + } + } + } + + pub fn next_downstream_prime(&mut self) -> Prime{ + loop { + if self.size - (self.prev_prime as f64).log2() >= 0.5 || self.prev_prime < self.nth_root{ + self.check_next_prime = false; + panic!("prime list for downstream primes is exhausted (overlap with previous bit-size or prime < nth_root)") + }else{ + if is_prime(self.prev_prime){ + let prime = Prime::new_unchecked(self.next_prime); + self.prev_prime -= self.nth_root; + return prime + } + self.prev_prime -= self.nth_root; + } + } + } + + pub fn next_alternating_prime(&mut self) -> Prime{ + loop { + if !(self.check_next_prime || self.check_prev_prime){ + panic!("prime list for upstream and downstream prime is exhausted for the (overlap with previous/next bit-size or NthRoot > prime > 2^64)") + } + + if self.check_next_prime{ + if (self.next_prime as f64).log2() - self.size >= 0.5 || self.next_prime > 0xffff_ffff_ffff_ffff-self.nth_root{ + self.check_next_prime = false; + }else{ + if is_prime(self.next_prime){ + let prime = Prime::new_unchecked(self.next_prime); + self.next_prime += self.nth_root; + return prime + } + self.next_prime += self.nth_root; + } + } + + if self.check_prev_prime { + if self.size - (self.prev_prime as f64).log2() >= 0.5 || self.prev_prime < self.nth_root{ + self.check_prev_prime = false; + }else{ + if is_prime(self.prev_prime){ + let prime = Prime::new_unchecked(self.prev_prime); + self.prev_prime -= self.nth_root; + return prime + } + self.prev_prime -= self.nth_root; + } + } + } + } +} + + +#[cfg(test)] +mod test { + use crate::modulus::prime_generator; + + #[test] + fn prime_generation() { + let nth_root: u64 = 1<<16 ; + let mut g: prime_generator::NTTFriendlyPrimesGenerator = prime_generator::NTTFriendlyPrimesGenerator::new(30, nth_root); + + let primes = g.next_alternating_primes(10); + println!("{:?}", primes); + for prime in primes.iter(){ + assert!(prime.q() % nth_root == 1); + } + } +} \ No newline at end of file