#![no_std]
//! This crate implements primality tests for the primitive integer types up to `u64`.
//! It is based on the rabin-miller primality test (except for `u8`, which uses a lookup table).
//! While the rabin-miller test is probabilistic, this crate uses sets of bases that are known to produce correct results for all values in the ranges of the respective integer types.
//! There is no such set known for `u128`, which is why `u128` is not supported.
//! ```
//! # use small_is_prime::{is_prime_u64,IsPrime};
//! assert!(!is_prime_u64(1000000000000));
//! assert!(1000000000039u64.is_prime());
//! ```

use num_modular::{Montgomery, Reducer};

const PRIME_BIT_MASK: u64 = 1 << 2
    | 1 << 3
    | 1 << 5
    | 1 << 7
    | 1 << 11
    | 1 << 13
    | 1 << 17
    | 1 << 19
    | 1 << 23
    | 1 << 29
    | 1 << 31
    | 1 << 37
    | 1 << 41
    | 1 << 43
    | 1 << 47
    | 1 << 53
    | 1 << 59
    | 1 << 61;

macro_rules! bit_lookup_test {
    ($scrutinee:expr,$T:ty,$bit_limit:expr) => {
        if $scrutinee < $bit_limit {
            const MASK: $T =
                (PRIME_BIT_MASK as $T) & <$T>::unbounded_shl(1, $bit_limit).wrapping_sub(1);
            return ((MASK as $T) >> ($scrutinee as u32)) & 1 != 0;
        }
    };
}

macro_rules! direct_test{
    ($scrutinee:expr$(,$p:expr)*)=>{
        if $($scrutinee % $p == 0)||* {
            return false
        }
    }
}

macro_rules! bases {
    (64) => {
        [2, 325, 9375, 28178, 450775, 9780504, 1795265022]
    };
    (32) => {
        [2, 7, 61]
    };
    (16) => {
        [2, 3]
    };
}

macro_rules! impl_for_primitive {
    ($T:ty as $As:ty,$f:ident) => {
        impl Sealed for $T {}
        impl IsPrime for $T {
            fn is_prime(self) -> bool {
                $f(self as $As)
            }
        }
    };
}

macro_rules! define_is_prime {
    ($name:ident,$T:ty,$bases:expr) => {
        pub fn $name(n: $T) -> bool {
            bit_lookup_test!(n, u32, 32);
            #[cfg(not(test))]
            direct_test!(n, 2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31);
            #[cfg(test)]
            direct_test!(n, 2);
            let s = (n - 1).trailing_zeros();
            let d = (n - 1) >> s;
            fn mod_pow(one: $T, mut base: $T, mut exp: $T, m: Montgomery<$T>) -> $T {
                let mut result = one;
                while exp > 0 {
                    if exp % 2 != 0 {
                        result = m.mul(&result, &base);
                    }
                    base = m.sqr(base);
                    exp /= 2;
                }
                result
            }
            let nm = Montgomery::<$T>::new(n);
            let onem = nm.transform(1);
            let neg_onem = nm.transform(n - 1);
            'next_base: for &a in &$bases {
                if a >= n {
                    continue;
                }
                let mut xm = mod_pow(onem, nm.transform(a), d, nm);
                if xm == onem || xm == neg_onem {
                    continue;
                }
                for _ in 1..s {
                    xm = nm.sqr(xm);
                    if xm == neg_onem {
                        continue 'next_base;
                    }
                    if xm == onem {
                        return false;
                    }
                }
                return false;
            }
            true
        }
        impl_for_primitive!($T as $T, $name);
    };
}

pub fn is_prime_u8(n: u8) -> bool {
    static LOOKUP: [u32; 8] = [
        0xa08a28ac, 0x28208a20, 0x2088288, 0x800228a2, 0x20a00a08, 0x80282088, 0x800800a2,
        0x8028228,
    ];
    (LOOKUP[n as usize / 32] >> (n as u32 % 32) & 1) != 0
}

trait Sealed {}
#[allow(private_bounds)]
pub trait IsPrime: Sealed {
    #[allow(clippy::wrong_self_convention)]
    fn is_prime(self) -> bool;
}

impl_for_primitive!(u8 as u8, is_prime_u8);
define_is_prime!(is_prime_u16, u16, bases!(16));
define_is_prime!(is_prime_u32, u32, bases!(32));
define_is_prime!(is_prime_u64, u64, bases!(64));
#[cfg(target_pointer_width = "32")]
impl_for_primitive!(usize as u32, is_prime_u32);
#[cfg(target_pointer_width = "64")]
impl_for_primitive!(usize as u64, is_prime_u64);
