prime-formula 0.3.1

//! Fast primality tests for large integers using Miller-Rabin tests
//!
//! Algorithm has been optimized to avoid overflow whilst keeping
//! everything 128-bit for speed.
//
// Copyright: Adam Cottrell (cottrela@gmail.com)

use rand::Rng;

// For u64 bit values, these are known to give deterministic results
const DEFAULT_ROOTS: [u8; 12] = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37];

fn decompose(n: &u128) -> (u128, u64) {
    let mut d = n - 1;
    let mut s = 0;

    // While d is even
    while d & 1 == 0 {
        d >>= 1;
        s += 1;
    }
    if d == 0 {
        d = 1;
    }
    (d, s)
}

fn is_witness(a: u128, d: u128, s: u64, n: &u128) -> bool {
    let mut x = mod_exp(a, d, n);

    if x == 1 || x == n - 1 {
        return false;
    }
    for _ in 0..s - 1 {
        x = mod_exp(x, 2, n);
        if x == n - 1 {
            return false;
        }
    }
    true
}

/// Convenience function for checking primality
pub fn is_prime(n: &u128, k: u32, deterministic: bool) -> bool {
    if deterministic {
        is_deterministic_prime(n, k)
    } else {
        is_likely_prime(n, k)
    }
}

/// Deterministic primality test
///
/// Guaranteed for any prime candidate in the range
/// 11 to 18,446,744,073,709,551,615
///
/// If you are looking for primality checks for primes < 11
/// then do not use this function!!
pub fn is_deterministic_prime(candidate: &u128, num_checks: u32) -> bool {
    let (d, s) = decompose(candidate);

    for i in 0..num_checks {
        let a = DEFAULT_ROOTS[i as usize] as u128;
        if a >= *candidate {
            continue;
        }
        if is_witness(a, d, s, candidate) {
            return false;
        }
    }
    true
}

/// Non-deterministic primality test
///
/// Performs Miller-Rabin algorithm, but with a random witness in the range
/// 2..n-2 where n is the candidate number.
///
/// Note that this function is optimized to work on groups of large primes
/// but will by design not check primes less than 11 (which is the first root
/// of prime)
pub fn is_likely_prime(candidate: &u128, num_checks: u32) -> bool {
    let (d, s) = decompose(candidate);
    let mut rng = rand::thread_rng();

    for _ in 0..num_checks {
        let a = rng.gen_range(2..candidate - 2);
        if is_witness(a, d, s, candidate) {
            return false;
        }
    }
    true
}

fn add_mod(mut a: u128, mut b: u128, m: &u128) -> u128 {
    // Equivalent of a + b % m, but with overflow protection
    a %= m;
    b %= m;
    let r = m - a;
    if b < r { a + b } else { b - r }
}

fn mul_mod(mut a: u128, mut b: u128, m: &u128) -> u128 {
    // Equivalent of a * b % m
    //
    // Code uses the Russian peasant algorithm to avoid
    // using 2xN sized container, and uses mod identities
    // where possible to reduce work done

    // ab mod c = (a mod c x b mod c) mod c
    a %= m;
    b %= m;

    // Swap a/b to ensure b is always smallest
    // this reduces big-O on the while loop
    if b > a {
        std::mem::swap(&mut a, &mut b);
        //         let tmp = b;
        //         b = a;
        //         a = tmp;
    }

    // Using the russian peasant algorithm to avoid
    // overflows on the multiply
    let mut res = 0;
    while b != 0 {
        // If b is odd add the residue
        if b & 1 == 1 {
            res = add_mod(res, a, m);
        }

        // Multiple by 2 - technically could overflow
        // but with no overflow checks
        a = a.checked_shl(1).unwrap();
        a %= m;

        // Divide by 2
        b >>= 1;
    }

    res % m
}

fn mod_exp(mut base: u128, mut exponent: u128, modulus: &u128) -> u128 {
    let mut result = 1;
    base %= modulus;
    while exponent != 0 {
        if exponent & 1 == 1 {
            result = mul_mod(result, base, modulus);
        }
        exponent >>= 1;
        base = mul_mod(base, base, modulus);
    }
    result
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_decompose() {
        assert_eq!(decompose(&13), (3, 2)); // 12 = 3 * 2^2
        assert_eq!(decompose(&9), (1, 3)); // 8 = 1 * 2^3
        assert_eq!(decompose(&17), (1, 4)); // 16 = 1 * 2^4
        assert_eq!(decompose(&7), (3, 1)); // 6 = 3 * 2^1
    }

    #[test]
    fn test_mod_exp() {
        assert_eq!(mod_exp(2, 3, &5), 3); // 2^3 % 5 = 3
        assert_eq!(mod_exp(3, 2, &4), 1); // 3^2 % 4 = 1
        assert_eq!(mod_exp(5, 0, &7), 1); // x^0 = 1
        assert_eq!(mod_exp(10, 3, &1000), 0); // 10^3 % 1000 = 0
    }

    #[test]
    fn test_modular_arithmetic() {
        // add_mod tests
        assert_eq!(add_mod(u128::MAX, 1, &10), 6); // (MAX % 10 + 1) % 10
        assert_eq!(add_mod(7, 8, &5), 0); // (7+8) % 5 = 15%5=0

        // mul_mod tests
        assert_eq!(mul_mod(2, 3, &5), 1); // 2*3 %5=1
        assert_eq!(mul_mod(u128::MAX, 2, &100), (u128::MAX % 100 * 2) % 100);
    }

    #[test]
    fn test_deterministic_primes() {
        // Known primes
        assert!(is_deterministic_prime(&9973, 5)); // Prime
        assert!(is_deterministic_prime(&15485863, 5)); // 1,000,000th prime

        // Known composites
        assert!(!is_deterministic_prime(&990007, 5)); // 997 * 991
        assert!(!is_deterministic_prime(&561, 5)); // Carmichael number
    }

    #[test]
    fn test_probabilistic_primes() {
        // Test with 10 iterations for reliability
        assert!(is_likely_prime(&7919, 10)); // Prime
        assert!(!is_likely_prime(&7921, 10)); // 89^2 = 7921

        // Large prime
        assert!(is_likely_prime(&0xFFFFFFFFFFFFFFC5u128, 10));
    }

    #[test]
    fn test_interface_functions() {
        // is_prime wrapper
        assert!(is_prime(&13, 5, true)); // Deterministic
        assert!(is_prime(&13, 5, false)); // Probabilistic
    }

    #[test]
    fn test_carmichael_numbers() {
        // These should all be detected as composite
        let carmichaels = [561, 1105, 1729, 2465, 2821];
        for n in carmichaels {
            assert!(!is_deterministic_prime(&n, 12)); // Test with all bases
        }
    }

    #[test]
    fn test_large_numbers() {
        // Mersenne prime (2^61 - 1)
        let mersenne61 = (1u128 << 61) - 1;
        assert!(is_deterministic_prime(&mersenne61, 5));

        // Large composite (m61 ^ 2)
        let large_composite = mersenne61 * mersenne61;
        assert!(!is_deterministic_prime(&large_composite, 5));
    }
}