composite_modulus_proofs 0.1.0

use crate::math::small_primes::SMALL_PRIMES;
use core::ops::{Mul, Sub};
use crypto_bigint::{BoxedUint, Odd, Uint};
use crypto_primes::is_prime_with_rng;
use rand_core::CryptoRngCore;

/// Check if `n = p^k` where `p` is an odd prime. Taken from https://cstheory.stackexchange.com/a/36096
// TODO: This can likely be improved by calculating n-th roots
pub fn is_prime_power<R: CryptoRngCore, const LIMBS: usize>(
    rng: &mut R,
    n: &Odd<Uint<LIMBS>>,
) -> Option<(Odd<Uint<LIMBS>>, u32)> {
    let lg_n = n.as_ref().bits_vartime();
    for b in 2..lg_n {
        let mut a_min = Uint::ONE;
        // 1 << ceil(lg_n/b)
        let mut a_max = Uint::ONE.shl_vartime((lg_n + b + 1) / b);
        while a_min.cmp_vartime(&a_max.sub(Uint::ONE)).is_lt() {
            let a_mid = (a_min + a_max).shr(1);
            let res = exp_with_square_and_multiply(a_mid, b);
            let n_box = BoxedUint::zero_with_precision(res.bits_precision()) + BoxedUint::from(n);
            if res.cmp_vartime(&n_box).is_gt() {
                a_max = a_mid;
            } else if res.cmp_vartime(&n_box).is_lt() {
                a_min = a_mid;
            } else {
                if is_prime_with_rng(rng, &a_mid) {
                    return Some((a_mid.to_odd().unwrap(), b));
                } else {
                    break;
                }
            }
        }
    }
    None
}

/// Checks if `n` is divisible by any prime in `[2, max_divisor]` and returns the first such prime it finds. Else returns None.
pub fn is_divisible_by_prime_upto<const LIMBS: usize>(
    n: &Odd<Uint<LIMBS>>,
    max_divisor: u32,
) -> Option<u32> {
    for i in 3..=max_divisor {
        if is_prime(i) {
            let rem = n.rem_vartime(&Uint::from(i).to_nz().unwrap());
            if rem == Uint::ZERO {
                return Some(i);
            }
        }
    }
    None
}

/// From https://stackoverflow.com/a/55790712
pub fn is_prime(n: u32) -> bool {
    // If its found in the sorted list of hardcoded primes
    let last_small_prime = *SMALL_PRIMES.last().unwrap();
    if n <= last_small_prime {
        return SMALL_PRIMES.binary_search(&n).is_ok();
    }

    use micromath::F32;
    let limit = F32(n as f32).sqrt().ceil().0 as u32;
    for i in 2..limit {
        if n % i == 0 {
            return false;
        }
    }
    true
}

pub fn exp_with_square_and_multiply<const BASE_LIMBS: usize>(
    base: Uint<BASE_LIMBS>,
    mut exp: u32,
) -> BoxedUint {
    let mut res = BoxedUint::one();
    let mut base = BoxedUint::from(base);
    while exp != 0 {
        if exp % 2 == 1 {
            res = res.mul(&base);
        }
        base = base.square();
        exp = exp >> 1;
    }
    res
}

#[cfg(test)]
mod tests {
    use super::*;
    use crypto_bigint::{U128, U64};
    use crypto_primes::{generate_prime_with_rng, is_prime_with_rng};
    use rand_core::OsRng;
    use std::time::Instant;

    #[test]
    fn check_prime() {
        let mut rng = OsRng::default();
        assert!(is_prime(13));
        assert!(is_prime(19));
        assert!(!is_prime(20));
        assert!(!is_prime(21));

        let last_small_prime = *SMALL_PRIMES.last().unwrap();
        assert!(is_prime(last_small_prime));
        assert!(!is_prime(104728));

        for _ in 0..100 {
            let p: U64 = generate_prime_with_rng(&mut rng, u32::BITS);
            assert!(is_prime_with_rng(&mut rng, &p));
            let p_u32 = p.to_words()[0] as u32;
            assert!(is_prime(p_u32), "{} {}", p, p_u32);
        }

        let p: U64 = generate_prime_with_rng(&mut rng, u32::BITS);
        let q: U64 = generate_prime_with_rng(&mut rng, u32::BITS);
        let n = p.mul(&q).resize::<{ U128::LIMBS }>().to_odd().unwrap();
        let p_u32 = p.to_words()[0] as u32;
        let q_u32 = q.to_words()[0] as u32;
        for i in 3..=100 {
            if i < p_u32 && i < q_u32 && (i % 2 == 1) {
                assert!(is_divisible_by_prime_upto(&n, i).is_none());
                let n_ = n.mul(U128::from(i)).to_odd().unwrap();
                assert!(is_divisible_by_prime_upto(&n_, i).is_some());
            }
        }
    }

    #[test]
    fn check_prime_power() {
        let mut rng = OsRng::default();
        assert!(is_prime_power(&mut rng, &U64::from(15_u32).to_odd().unwrap()).is_none());
        assert!(is_prime_power(&mut rng, &U64::from(225_u32).to_odd().unwrap()).is_none());
        assert_eq!(
            is_prime_power(&mut rng, &U64::from(25_u32).to_odd().unwrap()).unwrap(),
            (Uint::from(5_u32).to_odd().unwrap(), 2)
        );

        macro_rules! check {
            ( $num_iters:expr, $prime_type:ident ) => {
                let start = Instant::now();
                for _ in 0..100 {
                    let p: Odd<$prime_type> =
                        generate_prime_with_rng::<$prime_type>(&mut rng, $prime_type::BITS)
                            .to_odd()
                            .unwrap();
                    let q: Odd<$prime_type> =
                        generate_prime_with_rng::<$prime_type>(&mut rng, $prime_type::BITS)
                            .to_odd()
                            .unwrap();
                    let r: Odd<$prime_type> =
                        generate_prime_with_rng::<$prime_type>(&mut rng, $prime_type::BITS)
                            .to_odd()
                            .unwrap();

                    let n = p.widening_mul(&q).to_odd().unwrap();
                    assert!(is_prime_power(&mut rng, &n).is_none(), "{} {} {}", n, p, q);

                    let p_sqr = p.widening_mul(&p).to_odd().unwrap();
                    let q_sqr = q.widening_mul(&q).to_odd().unwrap();
                    assert_eq!(
                        is_prime_power(&mut rng, &p_sqr).unwrap(),
                        (p.resize().to_odd().unwrap(), 2),
                        "{} {}",
                        p_sqr,
                        p
                    );
                    assert_eq!(
                        is_prime_power(&mut rng, &q_sqr).unwrap(),
                        (q.resize().to_odd().unwrap(), 2),
                        "{} {}",
                        q_sqr,
                        q
                    );

                    let pqr = n.widening_mul(&r).to_odd().unwrap();
                    assert!(
                        is_prime_power(&mut rng, &pqr).is_none(),
                        "{} {} {} {}",
                        pqr,
                        p,
                        q,
                        r
                    );

                    let p_sqr_q = p_sqr.widening_mul(&q).to_odd().unwrap();
                    assert!(
                        is_prime_power(&mut rng, &p_sqr_q).is_none(),
                        "{} {} {}",
                        p_sqr_q,
                        p,
                        q
                    );

                    let p_sqr_q_sqr = p_sqr.widening_mul(&q_sqr).to_odd().unwrap();
                    assert!(
                        is_prime_power(&mut rng, &p_sqr_q_sqr).is_none(),
                        "{} {} {}",
                        p_sqr_q_sqr,
                        p,
                        q
                    );

                    let p_cube = p_sqr.widening_mul(&p).to_odd().unwrap();
                    let q_cube = q_sqr.widening_mul(&q).to_odd().unwrap();
                    assert_eq!(
                        is_prime_power(&mut rng, &p_cube).unwrap(),
                        (p.resize().to_odd().unwrap(), 3),
                        "{} {}",
                        p_cube,
                        p
                    );
                    assert_eq!(
                        is_prime_power(&mut rng, &q_cube).unwrap(),
                        (q.resize().to_odd().unwrap(), 3),
                        "{} {}",
                        q_cube,
                        q
                    );
                }
                println!(
                    "{} checks for {}-bit prime done in: {:?}",
                    $num_iters,
                    $prime_type::BITS,
                    start.elapsed()
                );
            };
        }

        check!(100, U64);
        check!(100, U128);
    }
}