Added Barrett & fixed Montgomery, added tests

Cargo.lock

@@ -2,13 +2,49 @@
 # It is not intended for manual editing.
 version = 4
 
+[[package]]
+name = "autocfg"
+version = "1.4.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ace50bade8e6234aa140d9a2f552bbee1db4d353f69b8217bc503490fc1a9f26"
+
 [[package]]
 name = "math"
 version = "0.1.0"
 dependencies = [
+ "num-bigint",
+ "num-traits",
  "primality-test",
 ]
 
+[[package]]
+name = "num-bigint"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9"
+dependencies = [
+ "num-integer",
+ "num-traits",
+]
+
+[[package]]
+name = "num-integer"
+version = "0.1.46"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "7969661fd2958a5cb096e56c8e1ad0444ac2bbcd0061bd28660485a44879858f"
+dependencies = [
+ "num-traits",
+]
+
+[[package]]
+name = "num-traits"
+version = "0.2.19"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841"
+dependencies = [
+ "autocfg",
+]
+
 [[package]]
 name = "primality-test"
 version = "0.3.0"

Cargo.toml

@@ -5,3 +5,5 @@ edition = "2021"
 
 [dependencies]
 primality-test = "0.3.0"
+num-bigint = "0.4.6"
+num-traits = "0.2.19"

src/lib.rs

@@ -1,3 +1,4 @@
 #![feature(bigint_helper_methods)]
+#![feature(test)]
 
 pub mod modulus;

src/modulus/barrett.rs

@@ -1,19 +1,28 @@
+use num_bigint::BigUint;
+use num_traits::cast::ToPrimitive;
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct BarrettPrecomp<O>(O, O);
 
 impl<O> BarrettPrecomp<O>{
 
-    #[inline(always)]
-    pub fn new(a:O, b: O) -> Self{
-        Self(a, b)
-    }
-
     #[inline(always)]
     pub fn value_hi(&self) -> &O{
-        &self.0
+        &self.1
     }
 
     #[inline(always)]
     pub fn value_lo(&self) -> &O{
-        &self.1
+        &self.0
+    }
+}
+
+impl BarrettPrecomp<u64>{
+    pub fn new(q: u64) -> BarrettPrecomp<u64> {
+        let mut big_r = BigUint::parse_bytes(b"100000000000000000000000000000000", 16).unwrap();
+        big_r = big_r / BigUint::from(q);
+        let lo = (&big_r & BigUint::from(u64::MAX)).to_u64().unwrap();
+        let hi = (big_r >> 64u64).to_u64().unwrap();
+        Self(lo, hi)
     }
 }

src/modulus/montgomery.rs

@@ -1,8 +1,11 @@
 use crate::modulus::barrett::BarrettPrecomp;
 use crate::modulus::ReduceOnce;
 
+extern crate test;
+
 /// Montgomery is a generic struct storing
 /// an element in the Montgomery domain.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct Montgomery<O>(O);
 
 /// Implements helper methods on the struct Montgomery<O>.
@@ -25,6 +28,7 @@ impl<O> Montgomery<O>{
 
 /// MontgomeryPrecomp is a generic struct storing 
 /// precomputations for Montgomery arithmetic.
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct MontgomeryPrecomp<O>{
     q: O,
     q_barrett: BarrettPrecomp<O>,
@@ -38,15 +42,15 @@ impl MontgomeryPrecomp<u64>{
     /// Returns an new instance of MontgomeryPrecomp<u64>.
     /// This method will fail if gcd(q, 2^64) != 1.
     #[inline(always)]
-    fn new(&self, q: u64) -> MontgomeryPrecomp<u64>{
+    fn new(q: u64) -> MontgomeryPrecomp<u64>{
         assert!(q & 1 != 0, "Invalid argument: gcd(q={}, radix=2^64) != 1", q);
-        let mut r: u64 = 1;
+        let mut q_inv: u64 = 1;
         let mut q_pow = q;
         for _i in 0..63{
-            r = r.wrapping_mul(r);
+            q_inv = q_inv.wrapping_mul(q_pow);
-            q_pow = q_pow.wrapping_mul(q_pow)
+            q_pow = q_pow.wrapping_mul(q_pow);
         }
-        Self{ q: q, q_barrett: BarrettPrecomp::new(q, q), q_inv: q_pow}
+        Self{ q: q, q_barrett: BarrettPrecomp::new(q), q_inv: q_inv}
     }
 
     /// Returns (lhs<<64)%q as a Montgomery<u64>.
@@ -75,7 +79,7 @@ impl MontgomeryPrecomp<u64>{
     /// Assigns (lhs<<64)%q in range [0, 2q-1] to rhs.
     #[inline(always)]
     fn prepare_lazy_assign(&self, lhs: u64, rhs: &mut Montgomery<u64>){
-        let (mhi, _) = lhs.widening_mul(*self.q_barrett.value_lo());
+        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());
     }
 
@@ -105,9 +109,9 @@ impl MontgomeryPrecomp<u64>{
     /// Assigns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1] to rhs.
     #[inline(always)]
     fn mul_external_lazy_assign(&self, lhs: Montgomery<u64>, rhs: &mut u64){
-        let (mhi, mlo) = lhs.value().widening_mul(*rhs);
+        let (mlo, mhi) = lhs.value().widening_mul(*rhs);
-        let (hhi, _) = self.q.widening_mul(mlo * self.q_inv);
+        let (_, hhi) = self.q.widening_mul(mlo.wrapping_mul(self.q_inv));
-        *rhs = mhi - hhi + self.q
+        *rhs = mhi.wrapping_add(self.q - hhi)
     }
 
     /// Returns lhs * rhs * (2^{64})^-1 mod q in range [0, 2q-1].
@@ -133,4 +137,32 @@ impl MontgomeryPrecomp<u64>{
     fn mul_internal_lazy_assign(&self, lhs: Montgomery<u64>, rhs: &mut Montgomery<u64>){
         self.mul_external_lazy_assign(lhs, rhs.value_mut());
     }
+}
+
+
+#[cfg(test)]
+mod tests {
+    use crate::modulus::montgomery;
+    use super::*;
+    use test::Bencher;
+
+    #[test]
+    fn test_mul_external() {
+        let q: u64 = 0x1fffffffffe00001;
+    	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);
+    }
+
+    #[bench]
+    fn bench_mul_external(b: &mut Bencher){
+        let q: u64 = 0x1fffffffffe00001;
+    	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));
+    }
 }

src/modulus/prime_generation.rs

@@ -19,7 +19,7 @@ impl NTTFriendlyPrimesGenerator {
 		let next_prime = (1<<bit_size) + 1;
 		let mut prev_prime = next_prime;
 
-		if next_prime > 0xffff_ffff_ffff_ffff-nth_root{
+		if next_prime > nth_root.wrapping_neg(){
 			check_next_prime = false;
 		}