pmath 0.1.0

//! Functions related to prime numbers.

use crate::{gcd, newtons_method};
use num_traits::{ConstOne, ConstZero, PrimInt, ToPrimitive};

#[cfg_attr(doc, katexit::katexit)]
/// A prime-counting function.
///
/// Estimates the number of primes less than or equal to $x$.
/// Uses the prime number theorem which states that the number of primes
/// less than or equal to $x$ is approximately $x / \\ln{(x)}$.
/// It is exact for $x < 11$ and guaranteed to be an underestimate for $x >= 11$.
/// # Arguments
/// * `x` - The integer to estimate the number of primes less than or equal to.
/// # Returns
/// * The estimated number of primes less than or equal to `x`.
/// # Panics
/// * If `x` is negative.
/// * If `x` cannot be converted to [f64].
/// # Example
/// ```
/// use pmath::primes::pcf;
///
/// // 2, 3, 5, 7
/// assert_eq!(pcf(7), 4.0);
///
/// // number of primes less than or equal to 100: 25
/// assert!(pcf(100u8) < 25.0);
/// ```
pub fn pcf<T>(x: T) -> f64
where
    T: ToPrimitive,
{
    let x = x.to_f64().expect("Cannot convert x to f64.");
    if x < 0.0 {
        panic!("x must be non-negative.");
    } else if x < 2.0 {
        0.0
    } else if x < 3.0 {
        1.0
    } else if x < 5.0 {
        2.0
    } else if x < 7.0 {
        3.0
    } else if x < 11.0 {
        4.0
    } else {
        x / x.ln()
    }
}

#[cfg_attr(doc, katexit::katexit)]
/// An inverse prime-counting function.
///
/// Estimates the number for which the prime-counting function is approximately $n$.
/// Uses the inverse of the prime number theorem.
/// The answer is approximated using Newton's method.
/// For $n <= 3$ it is calculated exactly using $\\lfloor n \\rfloor$ and
/// guaranteed to be an overestimate for $n > 3$.
/// # Arguments
/// * `n` - The integer to estimate the inverse of the prime-counting function for.
/// # Returns
/// * The estimated inverse of the prime-counting function for `n`.
/// # Panics
/// * If `n` is negative.
/// * If `n` cannot be converted to [f64].
/// # Example
/// ```
/// use pmath::primes::apcf;
///
/// assert_eq!(apcf(2), 3.0);
/// assert!(apcf(25) > 100.0);
/// ```
pub fn apcf<T>(n: T) -> f64
where
    T: ToPrimitive,
{
    let mut n = n.to_f64().expect("Cannot convert n to f64.");
    if n < 0.0 {
        panic!("n must be non-negative.");
    } else if n <= 3.0 {
        n = n.floor();
        match n {
            0.0 => 0.0,
            1.0 => 2.0,
            2.0 => 3.0,
            3.0 => 5.0,
            _ => unreachable!(),
        }
    } else {
        let x0 = n + 1.0;
        let precision = 1e-10;
        let function = |x: f64| n * x.ln() - x;
        let derivative = |x: f64| n / x - 1.0;
        newtons_method(x0, precision, function, derivative).unwrap()
    }
}

/// Check if two integers are coprime.
/// # Arguments
/// * `a` - The first integer.
/// * `b` - The second integer.
/// # Returns
/// * Whether the two integers are coprime.
/// # Panics
/// * If either of the integers is negative.
/// # Example
/// ```
/// use pmath::primes::coprime;
///
/// assert!(coprime(7, 20)); // 7 and 20 are coprime
/// assert!(!coprime(12, 18)); // 12 and 18 are not coprime
/// assert!(coprime(15, 28)); // 15 and 28 are coprime
/// assert!(!coprime(10, 25)); // 10 and 25 are not coprime
/// assert!(coprime(1, 1)); // 1 and 1 are coprime
/// assert!(coprime(1, 2)); // 1 and 2 are coprime
/// assert!(coprime(2, 3)); // 2 and 3 are coprime
/// assert!(!coprime(2, 4)); // 2 and 4 are not coprime
/// ```
pub fn coprime<T>(a: T, b: T) -> bool
where
    T: PrimInt + ConstZero + ConstOne,
{
    gcd(a, b) == T::ONE
}

/// Check if an integer is prime.
/// # Arguments
/// * `n` - The integer to check the primality of.
/// # Returns
/// * `bool` - Whether the integer is prime.
/// * `T` - The smallest divisor if the integer is not prime, otherwise `1`.
/// # Panics
/// * If `n` is less than `2`.
/// * If `n` cannot be converted to [f64].
/// # Example
/// ```
/// use pmath::primes::is_prime;
///
/// // 7 is prime
/// assert_eq!(is_prime(7), (true, 1));
/// // 12 is not prime, the smallest divisor is 2
/// assert_eq!(is_prime(12), (false, 2));
/// ```
pub fn is_prime<T>(n: T) -> (bool, T)
where
    T: PrimInt + ConstZero + ConstOne,
{
    let t2 = T::from(2).unwrap();
    if n < t2 {
        panic!("n must be greater than or equal to 2.");
    }
    let t3 = T::from(3).unwrap();
    let t6 = T::from(6).unwrap();

    if n == t2 || n == t3 {
        (true, T::ONE)
    } else if n % t2 == T::ZERO {
        (false, t2)
    } else if n % t3 == T::ZERO {
        (false, t3)
    } else {
        let upper_bound =
            T::from(n.to_f64().expect("Cannot convert n to f64.").sqrt().floor()).unwrap();
        let mut i = T::from(5).unwrap();
        while i <= upper_bound {
            if n % i == T::ZERO {
                return (false, i);
            } else if n % (i + t2) == T::ZERO {
                return (false, i + t2);
            }
            i = i + t6; // increment by 6
        }

        (true, T::ONE)
    }
}

/// The sieve of Eratosthenes.
///
/// Finds all primes less than or equal to `n`.
/// # Arguments
/// * `n` - The number to find all primes less than or equal to.
/// # Returns
/// * All primes less than or equal to `n`.
/// # Panics
/// * If the sieve requires more elements than can be represented by [usize].
/// # Example
/// ```
/// use pmath::primes::sieve_of_eratosthenes;
///
/// // primes less than or equal to 10: 2, 3, 5, 7
/// assert_eq!(sieve_of_eratosthenes(10), vec![2, 3, 5, 7]);
/// ```
pub fn sieve_of_eratosthenes<T>(n: T) -> Vec<T>
where
    T: PrimInt + ConstOne,
{
    let t2 = T::from(2).unwrap();
    if n < t2 {
        Vec::with_capacity(0)
    } else if n == t2 {
        vec![t2]
    } else {
        let mut results = vec![t2];
        let mut sieve = vec![true; ((n - T::ONE) / t2).to_usize().expect("Sieve too large.")];

        let ind_to_val = |i: usize| ((i as u64) << 1) + 3; // calculate number value from index in sieve
        let val_to_ind = |v: u64| ((v - 3) >> 1) as usize; // calculate index in sieve from number value

        for prime_ind in 0..sieve.len() {
            if sieve[prime_ind] {
                // get prime number value
                let prime_val = ind_to_val(prime_ind);
                let prime_val_2 = prime_val << 1; // prime_val * 2, used to skip even numbers

                // check all multiples of prime number value and mark them as not prime
                // start checking at prime_val^2 (all smaller multiples have already been checked by smaller primes)
                let mut check_val = prime_val * prime_val;
                let mut check_ind = val_to_ind(check_val);
                if check_ind >= sieve.len() {
                    break;
                }

                while check_ind < sieve.len() {
                    sieve[check_ind] = false;
                    // we want check_val to always be odd, prime_val is always odd, so we can just add prime_val * 2
                    // (because if we added 2 odd numbers we would get an even number)
                    check_val += prime_val_2;
                    check_ind = val_to_ind(check_val);
                }
            }
        }

        // convert sieve indices that are true to their corresponding number values and add them to results
        results.extend(sieve.into_iter().enumerate().filter_map(|(i, prime)| {
            if prime {
                Some(T::from(ind_to_val(i)).unwrap())
            } else {
                None
            }
        }));

        // return results
        results
    }
}