feanor-math 3.5.18

use oorandom;

use crate::DEFAULT_PROBABILISTIC_REPETITIONS;
use crate::algorithms::eea::signed_gcd;
use crate::homomorphism::*;
use crate::integer::*;
use crate::ordered::OrderedRingStore;
use crate::primitive_int::*;
use crate::ring::*;
use crate::rings::zn::{ZnOperation, ZnRing, ZnRingStore, choose_zn_impl, zn_64};

struct CheckIsFieldMillerRabin {
    probability_param: usize,
}

impl ZnOperation for CheckIsFieldMillerRabin {
    type Output<'a>
        = bool
    where
        Self: 'a;

    fn call<'a, R>(self, ring: R) -> bool
    where
        R: 'a + ZnRingStore,
        R::Type: ZnRing,
    {
        is_prime_base(ring, self.probability_param)
    }
}

/// Miller-Rabin primality test.
///
/// If n is a prime, this returns true.
/// If n is not a prime, this returns false with probability greater or
/// equal than 1 - 4^(-k).
/// For very small primes, a lookup table may be used.
///
/// For details, see [`is_prime_base()`]
pub fn is_prime<I>(ZZ: I, n: &El<I>, k: usize) -> bool
where
    I: IntegerRingStore + HashableElRingStore,
    I::Type: IntegerRing + CanIsoFromTo<StaticRingBase<i128>>,
{
    assert!(ZZ.is_pos(n));
    if ZZ.is_zero(n) || ZZ.is_one(n) {
        false
    } else {
        let n_copy = ZZ.clone_el(n);
        choose_zn_impl(ZZ, n_copy, CheckIsFieldMillerRabin { probability_param: k })
    }
}

const SMALL_IS_COPRIME_TABLE: [bool; 30] = [
    false, true, false, false, false, false, false, true, false, false, false, true, false, true, false, false, false,
    true, false, true, false, false, false, true, false, false, false, false, false, true,
];

fn search_prime<I: IntegerRingStore>(ZZ: I, mut n: El<I>, delta: i64) -> Option<El<I>>
where
    I::Type: IntegerRing,
    zn_64::ZnBase: CanHomFrom<I::Type>,
{
    assert!(ZZ.is_pos(&n));

    let m = SMALL_IS_COPRIME_TABLE.len();
    let Zm = zn_64::Zn::new(m as u64);
    let mut n_mod_m = Zm.coerce(&ZZ, ZZ.clone_el(&n));
    let mut diff_to_n = 0;
    let Zi64_to_Zm = Zm.can_hom::<StaticRing<i64>>(&StaticRing::<i64>::RING).unwrap();
    let Zi64_to_ZZ = ZZ.can_hom(&StaticRing::<i64>::RING).unwrap();
    debug_assert!(ZZ.is_one(&signed_gcd(ZZ.clone_el(&n), Zi64_to_ZZ.map(delta), &ZZ)));
    let Zm_delta = Zi64_to_Zm.map(delta);
    let ZZ_m = Zi64_to_ZZ.map(m.try_into().unwrap());

    // we continue the main loop until we reach `n <= m`; at this point, the main loop is not
    // correct anymore, since being in `Z/mZ \ Z/mZ*` does not imply nonprimality anymore
    let mut remaining_steps = if ZZ.is_leq(&n, &ZZ_m) {
        0
    } else if delta < 0 {
        let max_steps = ZZ.ceil_div(
            ZZ.sub_ref(&n, &Zi64_to_ZZ.map(m.try_into().unwrap())),
            &Zi64_to_ZZ.map(-delta),
        );
        if ZZ.is_lt(&max_steps, &Zi64_to_ZZ.map(i64::MAX)) {
            int_cast(max_steps, StaticRing::<i64>::RING, &ZZ)
        } else {
            i64::MAX
        }
    } else {
        i64::MAX
    };

    while remaining_steps != 0 {
        while !SMALL_IS_COPRIME_TABLE[Zm.smallest_positive_lift(n_mod_m) as usize] && remaining_steps != 0 {
            diff_to_n += delta;
            Zm.add_assign(&mut n_mod_m, Zm_delta);
            remaining_steps -= 1;
        }
        ZZ.add_assign(&mut n, Zi64_to_ZZ.map(diff_to_n));
        if remaining_steps == 0 {
            break;
        }
        if is_prime(&ZZ, &n, DEFAULT_PROBABILISTIC_REPETITIONS) {
            return Some(n);
        } else {
            diff_to_n = delta;
            Zm.add_assign(&mut n_mod_m, Zm_delta);
            remaining_steps -= 1;
        }
    }
    let mut n = int_cast(n, StaticRing::<i64>::RING, &ZZ);
    assert!(n <= m.try_into().unwrap());
    while n > 0 {
        if is_prime(&StaticRing::<i64>::RING, &n, DEFAULT_PROBABILISTIC_REPETITIONS) {
            return Some(Zi64_to_ZZ.map(n));
        }
        n += delta;
    }
    return None;
}

/// Returns the largest prime smaller than the given integer.
#[stability::unstable(feature = "enable")]
pub fn prev_prime<I: IntegerRingStore>(ZZ: I, n: El<I>) -> Option<El<I>>
where
    I::Type: IntegerRing,
    zn_64::ZnBase: CanHomFrom<I::Type>,
{
    assert!(ZZ.is_pos(&n));
    if ZZ.is_one(&n) {
        return None;
    }
    let n_minus_one = ZZ.sub(n, ZZ.one());
    search_prime(ZZ, n_minus_one, -1)
}

/// Returns the smallest prime larger than the given integer.
#[stability::unstable(feature = "enable")]
pub fn next_prime<I: IntegerRingStore>(ZZ: I, n: El<I>) -> El<I>
where
    I::Type: IntegerRing,
    zn_64::ZnBase: CanHomFrom<I::Type>,
{
    assert!(ZZ.is_pos(&n));
    let n_plus_one = ZZ.add(n, ZZ.one());
    search_prime(ZZ, n_plus_one, 1).unwrap()
}

/// Miller-Rabin primality test.
///
/// If the characteristic `n` of the given ring is a prime, this returns true.
/// If n is not a prime, this returns false with probability greater or
/// equal than 1 - 4^(-k).
///
/// Complexity O(k log(n)^3)
///
/// # Randomness
///
/// Note that the randomness used for this function is derived only from
/// the input, hence it will always yield the same output on the same input.
/// Technically, it follows that the probability of a wrong output is greater
/// than 4^(-k) on some outputs (as it is either 0 or 1), but of course
/// this is not helpful. To be completely precise: If the seed of the used
/// PRNG would be random, then the probability of a wrong output is at
/// most 4^(-k).
pub fn is_prime_base<R>(Zn: R, k: usize) -> bool
where
    R: ZnRingStore,
    R::Type: ZnRing,
{
    let ZZ = Zn.integer_ring();
    let n = Zn.modulus();
    assert!(ZZ.is_pos(n));

    let mut rng = oorandom::Rand64::new(ZZ.default_hash(n) as u128);
    let mut n_minus_one = ZZ.sub_ref_fst(n, ZZ.one());
    let s = ZZ.abs_lowest_set_bit(&n_minus_one).unwrap();
    ZZ.euclidean_div_pow_2(&mut n_minus_one, s);
    let d = n_minus_one;

    for _i in 0..k {
        let a = Zn.random_element(|| rng.rand_u64());
        if Zn.is_zero(&a) {
            continue;
        }
        let mut current = Zn.pow_gen(a, &d, &ZZ);
        let mut miller_rabin_condition = Zn.is_one(&current);
        for _r in 0..s {
            miller_rabin_condition |= Zn.is_neg_one(&current);
            if miller_rabin_condition {
                break;
            }
            Zn.square(&mut current);
        }
        if Zn.is_zero(&current) || !miller_rabin_condition {
            return false;
        }
    }
    return true;
}

#[cfg(test)]
use crate::rings::rust_bigint::RustBigintRing;

#[test]
fn test_is_prime() {
    assert!(is_prime(StaticRing::<i128>::RING, &2, 5));
    assert!(is_prime(StaticRing::<i128>::RING, &3, 5));
    assert!(is_prime(StaticRing::<i128>::RING, &5, 5));
    assert!(is_prime(StaticRing::<i128>::RING, &7, 5));
    assert!(is_prime(StaticRing::<i128>::RING, &11, 5));
    assert!(is_prime(StaticRing::<i128>::RING, &22531, 5));
    assert!(is_prime(StaticRing::<i128>::RING, &417581, 5));
    assert!(is_prime(StaticRing::<i128>::RING, &68719476767, 5));

    assert!(!is_prime(StaticRing::<i128>::RING, &1, 5));
    assert!(!is_prime(StaticRing::<i128>::RING, &4, 5));
    assert!(!is_prime(StaticRing::<i128>::RING, &6, 5));
    assert!(!is_prime(StaticRing::<i128>::RING, &8, 5));
    assert!(!is_prime(StaticRing::<i128>::RING, &9, 5));
    assert!(!is_prime(StaticRing::<i128>::RING, &10, 5));
    assert!(!is_prime(StaticRing::<i128>::RING, &22532, 5));
    assert!(!is_prime(StaticRing::<i128>::RING, &347584, 5));

    assert!(is_prime(
        RustBigintRing::RING,
        &RustBigintRing::RING
            .get_ring()
            .parse("170141183460469231731687303715884105727", 10)
            .unwrap(),
        10
    ));
}

#[test]
fn test_prev_prime() {
    assert_eq!(Some(7), prev_prime(StaticRing::<i64>::RING, 11));
    assert_eq!(Some(7), prev_prime(StaticRing::<i64>::RING, 10));
    assert_eq!(Some(7), prev_prime(StaticRing::<i64>::RING, 9));
    assert_eq!(Some(7), prev_prime(StaticRing::<i64>::RING, 8));
    assert_eq!(Some(5), prev_prime(StaticRing::<i64>::RING, 7));
    assert_eq!(Some(5), prev_prime(StaticRing::<i64>::RING, 6));
    assert_eq!(Some(3), prev_prime(StaticRing::<i64>::RING, 5));
    assert_eq!(Some(3), prev_prime(StaticRing::<i64>::RING, 4));
    assert_eq!(Some(2), prev_prime(StaticRing::<i64>::RING, 3));
    assert_eq!(None, prev_prime(StaticRing::<i64>::RING, 2));
    assert_eq!(None, prev_prime(StaticRing::<i64>::RING, 1));

    let mut last_prime = 11;
    for i in 12..1000 {
        assert_eq!(Some(last_prime), prev_prime(StaticRing::<i64>::RING, i));
        if is_prime(StaticRing::<i64>::RING, &i, DEFAULT_PROBABILISTIC_REPETITIONS) {
            last_prime = i;
        }
    }
}

#[test]
fn test_next_prime() {
    let mut last_prime = 1009;
    for i in (2..1000).rev() {
        assert_eq!(last_prime, next_prime(StaticRing::<i64>::RING, i));
        if is_prime(StaticRing::<i64>::RING, &i, DEFAULT_PROBABILISTIC_REPETITIONS) {
            last_prime = i;
        }
    }
    assert_eq!(2, next_prime(StaticRing::<i64>::RING, 1));
}

#[test]
fn test_search_prime() {
    assert_eq!(None, search_prime(StaticRing::<i64>::RING, 1, -2));
    for (p, n) in [(3, 3), (5, 5), (7, 7), (7, 9), (11, 11)] {
        assert_eq!(Some(p), search_prime(StaticRing::<i64>::RING, n, -2));
    }

    assert_eq!(None, search_prime(StaticRing::<i64>::RING, 1, -3));
    assert_eq!(None, search_prime(StaticRing::<i64>::RING, 4, -3));
    for (p, n) in [
        (2, 2),
        (5, 5),
        (7, 7),
        (5, 8),
        (7, 10),
        (11, 11),
        (13, 13),
        (11, 14),
        (13, 16),
        (17, 17),
    ] {
        assert_eq!(Some(p), search_prime(StaticRing::<i64>::RING, n, -3));
    }
    assert_eq!(Some(359), search_prime(StaticRing::<i64>::RING, 380, -3));
}