factor: Faster modular arithmetic with the Montgomery transform

src/uu/factor/build.rs

@@ -20,56 +20,19 @@
 use std::env::{self, args};
 use std::fs::File;
 use std::io::Write;
-use std::num::Wrapping;
 use std::path::Path;
-use std::u64::MAX as MAX_U64;
 
 use self::sieve::Sieve;
 
 #[cfg(test)]
 use miller_rabin::is_prime;
 
-#[cfg(test)]
 #[path = "src/numeric.rs"]
 mod numeric;
+use numeric::inv_mod_u64;
 
 mod sieve;
 
-// extended Euclid algorithm
-// precondition: a does not divide 2^64
-fn inv_mod_u64(a: u64) -> Option<u64> {
-    let mut t = 0u64;
-    let mut newt = 1u64;
-    let mut r = 0u64;
-    let mut newr = a;
-
-    while newr != 0 {
-        let quot = if r == 0 {
-            // special case when we're just starting out
-            // This works because we know that
-            // a does not divide 2^64, so floor(2^64 / a) == floor((2^64-1) / a);
-            MAX_U64
-        } else {
-            r
-        } / newr;
-
-        let (tp, Wrapping(newtp)) = (newt, Wrapping(t) - (Wrapping(quot) * Wrapping(newt)));
-        t = tp;
-        newt = newtp;
-
-        let (rp, Wrapping(newrp)) = (newr, Wrapping(r) - (Wrapping(quot) * Wrapping(newr)));
-        r = rp;
-        newr = newrp;
-    }
-
-    if r > 1 {
-        // not invertible
-        return None;
-    }
-
-    Some(t)
-}
-
 #[cfg_attr(test, allow(dead_code))]
 fn main() {
     let out_dir = env::var("OUT_DIR").unwrap();
@@ -95,7 +58,7 @@ fn main() {
     let mut x = primes.next().unwrap();
     for next in primes {
         // format the table
-        let outstr = format!("({}, {}, {}),", x, inv_mod_u64(x).unwrap(), MAX_U64 / x);
+        let outstr = format!("({}, {}, {}),", x, inv_mod_u64(x), std::u64::MAX / x);
         if cols + outstr.len() > MAX_WIDTH {
             write!(file, "\n    {}", outstr).unwrap();
             cols = 4 + outstr.len();
@@ -116,18 +79,12 @@ fn main() {
 }
 
 #[test]
-fn test_inverter() {
-    let num = 10000;
-
-    let invs = Sieve::odd_primes().map(|x| inv_mod_u64(x).unwrap());
-    assert!(Sieve::odd_primes().zip(invs).take(num).all(|(x, y)| {
-        let Wrapping(z) = Wrapping(x) * Wrapping(y);
-        is_prime(x) && z == 1
-    }));
+fn test_generator_isprime() {
+    assert_eq!(Sieve::odd_primes.take(10_000).all(is_prime));
 }
 
 #[test]
-fn test_generator() {
+fn test_generator_10001() {
     let prime_10001 = Sieve::primes().skip(10_000).next();
     assert_eq!(prime_10001, Some(104_743));
 }

src/uu/factor/src/factor.rs

@@ -39,14 +39,21 @@ impl Factors {
     }
 
     fn add(&mut self, prime: u64, exp: u8) {
-        assert!(exp > 0);
+        debug_assert!(exp > 0);
         let n = *self.f.get(&prime).unwrap_or(&0);
         self.f.insert(prime, exp + n);
     }
 
     fn push(&mut self, prime: u64) {
         self.add(prime, 1)
     }
+
+    #[cfg(test)]
+    fn product(&self) -> u64 {
+        self.f
+            .iter()
+            .fold(1, |acc, (p, exp)| acc * p.pow(*exp as u32))
+    }
 }
 
 impl ops::MulAssign<Factors> for Factors {
@@ -132,3 +139,22 @@ pub fn uumain(args: Vec<String>) -> i32 {
     }
     0
 }
+
+#[cfg(test)]
+mod tests {
+    use super::factor;
+
+    #[test]
+    fn factor_recombines_small() {
+        assert!((1..10_000)
+            .map(|i| 2 * i + 1)
+            .all(|i| factor(i).product() == i));
+    }
+
+    #[test]
+    fn factor_recombines_overflowing() {
+        assert!((0..250)
+            .map(|i| 2 * i + 2u64.pow(32) + 1)
+            .all(|i| factor(i).product() == i));
+    }
+}

src/uu/factor/src/miller_rabin.rs

@@ -23,9 +23,10 @@ impl Result {
 // Deterministic Miller-Rabin primality-checking algorithm, adapted to extract
 // (some) dividers; it will fail to factor strong pseudoprimes.
 #[allow(clippy::many_single_char_names)]
-pub(crate) fn test<A: Arithmetic>(n: u64) -> Result {
+pub(crate) fn test<A: Arithmetic>(m: A) -> Result {
     use self::Result::*;
 
+    let n = m.modulus();
     if n < 2 {
         return Pseudoprime;
     }
@@ -37,36 +38,41 @@ pub(crate) fn test<A: Arithmetic>(n: u64) -> Result {
     let i = (n - 1).trailing_zeros();
     let r = (n - 1) >> i;
 
-    for a in BASIS.iter() {
-        let a = a % n;
-        if a == 0 {
+    let one = m.one();
+    let minus_one = m.minus_one();
+
+    for _a in BASIS.iter() {
+        let _a = _a % n;
+        if _a == 0 {
             break;
         }
 
+        let a = m.from_u64(_a);
+
         // x = a^r mod n
-        let mut x = A::pow(a, r, n);
+        let mut x = m.pow(a, r);
 
         {
             // y = ((x²)²...)² i times = x ^ (2ⁱ) = a ^ (r 2ⁱ) = x ^ (n - 1)
             let mut y = x;
             for _ in 0..i {
-                y = A::mul(y, y, n)
+                y = m.mul(y, y)
             }
-            if y != 1 {
+            if y != one {
                 return Pseudoprime;
             };
         }
 
-        if x == 1 || x == n - 1 {
+        if x == one || x == minus_one {
             break;
         }
 
         loop {
-            let y = A::mul(x, x, n);
-            if y == 1 {
-                return Composite(gcd(x - 1, n));
+            let y = m.mul(x, x);
+            if y == one {
+                return Composite(gcd(m.to_u64(x) - 1, m.modulus()));
             }
-            if y == n - 1 {
+            if y == minus_one {
                 // This basis element is not a witness of `n` being composite.
                 // Keep looking.
                 break;
@@ -81,10 +87,25 @@ pub(crate) fn test<A: Arithmetic>(n: u64) -> Result {
 // Used by build.rs' tests
 #[allow(dead_code)]
 pub(crate) fn is_prime(n: u64) -> bool {
-    if n < 1 << 63 {
-        test::<Small>(n)
-    } else {
-        test::<Big>(n)
+    test::<Montgomery>(Montgomery::new(n)).is_prime()
+}
+
+#[cfg(test)]
+mod tests {
+    use super::is_prime;
+    const LARGEST_U64_PRIME: u64 = 0xFFFFFFFFFFFFFFC5;
+
+    #[test]
+    fn largest_prime() {
+        assert!(is_prime(LARGEST_U64_PRIME));
+    }
+
+    #[test]
+    fn first_primes() {
+        use crate::table::{NEXT_PRIME, P_INVS_U64};
+        for (p, _, _) in P_INVS_U64.iter() {
+            assert!(is_prime(*p), "{} reported composite", p);
+        }
+        assert!(is_prime(NEXT_PRIME));
     }
-    .is_prime()
 }