alkahest-cas 3.12.0

//! V3-1 — Integer number theory via FLINT `fmpz`.
//!
//! Wraps proven primality, integer factorisation, totients, Jacobi symbols,
//! modular square roots (`fmpz_sqrtmod`), nth roots modulo primes when Coprime holds,
//! brute-force discrete logs, and quadratic Dirichlet characters (odd square-free conductor).

/// The classical arithmetic functions — the partition function, Bernoulli,
/// Euler, Stirling and harmonic numbers, Möbius μ, σ_k and sums of squares.
///
/// A child module rather than more entry points here, because its refusals
/// live on a separate `#[non_exhaustive]` enum: [`NumberTheoryError`] below is
/// exhaustive and in the stable surface, so a new variant on it would force a
/// major version bump. See [`arith::ArithmeticError`].
pub mod arith;

use crate::errors::AlkahestError;
use crate::flint::ffi;
use crate::flint::integer::FlintIntFactor;
use crate::flint::FlintInteger;
use rug::ops::RemRounding;
use rug::Complete;
use rug::Integer;
use std::cmp::Ordering;
use std::fmt;
use std::str::FromStr;

// ---------------------------------------------------------------------------
// NumberTheoryError
// ---------------------------------------------------------------------------

/// Failed integer number-theory primitive (`E-NT-*`).
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum NumberTheoryError {
    InvalidInput { msg: &'static str },
    Domain { msg: &'static str },
    NoSolution,
    CompositeModulus,
    UnsupportedNthRoot,
}

impl fmt::Display for NumberTheoryError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            NumberTheoryError::InvalidInput { msg } => write!(f, "{msg}"),
            NumberTheoryError::Domain { msg } => write!(f, "{msg}"),
            NumberTheoryError::NoSolution => {
                write!(f, "no discrete logarithm or modular root exists")
            }
            NumberTheoryError::CompositeModulus => write!(f, "operation requires a prime modulus"),
            NumberTheoryError::UnsupportedNthRoot => {
                write!(f, "nth root modulo p requires k=2 or gcd(k,p−1)=1")
            }
        }
    }
}

impl std::error::Error for NumberTheoryError {}

impl AlkahestError for NumberTheoryError {
    fn code(&self) -> &'static str {
        match self {
            NumberTheoryError::InvalidInput { .. } => "E-NT-001",
            NumberTheoryError::Domain { .. } => "E-NT-002",
            NumberTheoryError::NoSolution => "E-NT-003",
            NumberTheoryError::CompositeModulus => "E-NT-004",
            NumberTheoryError::UnsupportedNthRoot => "E-NT-005",
        }
    }

    fn remediation(&self) -> Option<&'static str> {
        match self {
            NumberTheoryError::InvalidInput { .. } => {
                Some("pass arbitrary-precision integers as decimal strings without spaces")
            }
            NumberTheoryError::Domain { .. } => {
                Some("check parity, positivity, and defined ranges")
            }
            NumberTheoryError::NoSolution => {
                Some("verify solvability: residue in ⟨base⟩, or quadratic residue for k=2")
            }
            NumberTheoryError::CompositeModulus => {
                Some("use a prime field modulus where the FLINT primitives apply")
            }
            NumberTheoryError::UnsupportedNthRoot => Some(
                "use sqrt (k=2) or primes with gcd(k,p−1)=1; Tonelli–Shanks chains are deferred",
            ),
        }
    }
}

fn parse_int(s: &str) -> Result<Integer, NumberTheoryError> {
    Integer::from_str(s.trim()).map_err(|_| NumberTheoryError::InvalidInput {
        msg: "invalid decimal integer string",
    })
}

fn parse_nonnegative(s: &str) -> Result<Integer, NumberTheoryError> {
    let z = parse_int(s)?;
    if z.cmp0() == Ordering::Less {
        Err(NumberTheoryError::Domain {
            msg: "expected a non-negative integer",
        })
    } else {
        Ok(z)
    }
}

/// Multiplicative inverse of `a` modulo `m` when `gcd(a,m)=1`.
fn mod_inverse(mut a: Integer, m: &Integer) -> Option<Integer> {
    if m.cmp0() != Ordering::Greater {
        return None;
    }
    if m == &Integer::from(1) {
        return Some(Integer::from(0));
    }
    a %= m;
    let (g, s, _) = a.extended_gcd(m.clone(), Integer::new());
    if g != 1 && g != -1 {
        return None;
    }
    let mut inv = if g == -1 { -s } else { s };
    inv %= m;
    if inv.cmp0() == Ordering::Less {
        inv += m;
    }
    Some(inv)
}

fn integer_is_odd(n: &Integer) -> bool {
    (n.clone() % Integer::from(2_u32)).cmp0() != Ordering::Equal
}

/// Positive integer parser (strictly \(> 0\) when required by the caller).
fn parse_positive(s: &str) -> Result<Integer, NumberTheoryError> {
    let z = parse_nonnegative(s)?;
    if z.is_zero() {
        Err(NumberTheoryError::Domain {
            msg: "expected a positive integer",
        })
    } else {
        Ok(z)
    }
}

/// Exact primality (`fmpz_is_prime`).
pub fn isprime(n: &str) -> Result<bool, NumberTheoryError> {
    let z = parse_int(n)?;
    if z.cmp0() != Ordering::Greater {
        return Err(NumberTheoryError::Domain {
            msg: "isprime expects a positive integer",
        });
    }
    if z < 2 {
        return Ok(false);
    }
    let fz = FlintInteger::from_rug(&z);
    let r = unsafe { ffi::fmpz_is_prime(fz.inner_ptr()) };
    Ok(r != 0)
}

/// Full factorisation: `(sign, list of (prime, exponent))` for \(\prod p^e\cdot \mathrm{sign}\).
///
/// The primes are listed in ascending order.
pub fn factorint(n: &str) -> Result<(i32, Vec<(String, u64)>), NumberTheoryError> {
    let z = parse_int(n)?;
    if z.is_zero() {
        return Err(NumberTheoryError::Domain {
            msg: "factorint(0) is undefined",
        });
    }
    let fz = FlintInteger::from_rug(&z);
    // FlintIntFactor is drop-safe — no manual fmpz_factor_clear needed.
    let mut fac = FlintIntFactor::new();
    fac.factor(&fz);
    // Ascending primes, whatever order FLINT's sieve or ECM found them in.
    let mut primes: Vec<(rug::Integer, u64)> = (0..fac.len())
        .map(|i| (fac.base_at(i).to_rug(), fac.exp_at(i)))
        .collect();
    primes.sort_by(|a, b| a.0.cmp(&b.0));
    let out = primes
        .into_iter()
        .map(|(p, e)| (p.to_string(), e))
        .collect();
    Ok((fac.sign(), out))
}

/// Composite cofactors at most this many bits are handed to FLINT's full
/// `fmpz_factor` even under a budget: its quadratic sieve finishes a 160-bit
/// (48-digit) semiprime in well under a second, and without it the only way
/// to split a balanced one is ECM luck.
const FACTOR_FULL_BITS: u64 = 160;

/// The `fmpz_factor_smooth` effort ladder, in bits of the factors sought.
/// Small steps, because a single step is one uninterruptible FLINT call and
/// the budget can only be looked at between them.
const FACTOR_SMOOTH_LADDER: [i64; 12] = [16, 20, 24, 28, 32, 36, 40, 44, 48, 52, 56, 60];

/// [`factorint`] that honours an active [`crate::budget`].
///
/// `fmpz_factor` is one uninterruptible FLINT call, and on a hard composite
/// it runs for as long as the number takes: `(2⁸⁹−1)(2¹⁰⁷−1)(2¹²⁷−1)`, three
/// Mersenne primes and no small factor, ran past twenty seconds under
/// `Budget(wall_ms=1000)` (audit C6). Under a budget this instead climbs a
/// ladder of bounded-effort `fmpz_factor_smooth` passes (trial division, then
/// ECM tuned for factors of the given size), checking the budget between
/// passes, and hands any composite cofactor of at most
/// 160 bits to the full factoriser. A composite that is
/// still unsplit when the budget runs out — or once the ladder is exhausted,
/// which means it has no factor ECM finds cheaply and is larger than the
/// sieve is handed — is an [`arith::ArithmeticError::Resource`] refusal carrying the
/// `E-BUDGET-*` trip (for an exhausted ladder, the budget's own error once it
/// fires; the ladder then keeps retrying the last rung until it does).
///
/// With **no** budget active this is exactly [`factorint`]: unbounded, as
/// documented. The overshoot past the wall-clock limit is at most one ladder
/// step (measured: well under a second for a 323-bit composite).
pub fn factorint_bounded(n: &str) -> Result<(i32, Vec<(String, u64)>), arith::ArithmeticError> {
    use crate::budget::BudgetTrip;
    if !crate::budget::is_active() {
        return Ok(factorint(n)?);
    }
    let z = parse_int(n)?;
    if z.is_zero() {
        return Err(NumberTheoryError::Domain {
            msg: "factorint(0) is undefined",
        }
        .into());
    }
    let trip = |e| arith::ArithmeticError::Resource(BudgetTrip::Budget(e));
    let sign = if z.cmp0() == Ordering::Less { -1 } else { 1 };

    // Primes found so far (with multiplicity), and composites still to split.
    let mut primes: Vec<(Integer, u64)> = Vec::new();
    let mut pending: Vec<(Integer, u64)> = vec![(z.abs(), 1)];
    let mut rung = 0usize;
    while let Some((c, e)) = pending.pop() {
        crate::budget::check().map_err(trip)?;
        if c == 1 {
            continue;
        }
        let fc = FlintInteger::from_rug(&c);
        let prime = unsafe { ffi::fmpz_is_probabprime(fc.inner_ptr()) } != 0
            && unsafe { ffi::fmpz_is_prime(fc.inner_ptr()) } != 0;
        if prime {
            primes.push((c, e));
            continue;
        }
        let mut fac = FlintIntFactor::new();
        if c.significant_bits() as u64 <= FACTOR_FULL_BITS {
            fac.factor(&fc);
        } else {
            let bits = FACTOR_SMOOTH_LADDER[rung.min(FACTOR_SMOOTH_LADDER.len() - 1)];
            rung += 1;
            fac.factor_smooth(&fc, bits, true);
        }
        let parts: Vec<(Integer, u64)> = (0..fac.len())
            .map(|i| (fac.base_at(i).to_rug(), fac.exp_at(i) * e))
            .collect();
        if parts.len() == 1 && parts[0].0 == c {
            // No progress on this pass: retry on the next rung.
            pending.push((c, e));
            continue;
        }
        pending.extend(parts);
    }

    // Merge equal primes (a prime can surface from two cofactors) and sort.
    primes.sort_by(|a, b| a.0.cmp(&b.0));
    let mut out: Vec<(String, u64)> = Vec::new();
    let mut last: Option<Integer> = None;
    for (p, e) in primes {
        if last.as_ref() == Some(&p) {
            if let Some(entry) = out.last_mut() {
                entry.1 += e;
            }
        } else {
            out.push((p.to_string(), e));
            last = Some(p);
        }
    }
    Ok((sign, out))
}

/// Next prime strictly after `n` (`fmpz_nextprime`).
pub fn nextprime(n: &str, proved: bool) -> Result<String, NumberTheoryError> {
    let z = parse_int(n)?;
    let fz = FlintInteger::from_rug(&z);
    let mut res = FlintInteger::new();
    unsafe {
        ffi::fmpz_nextprime(
            res.inner_mut_ptr(),
            fz.inner_ptr(),
            if proved { 1 } else { 0 },
        );
    }
    Ok(res.to_string())
}

/// Euler totient \(\varphi(n)\) (`fmpz_euler_phi`).
pub fn totient(n: &str) -> Result<String, NumberTheoryError> {
    let z = parse_positive(n)?;
    let fz = FlintInteger::from_rug(&z);
    let mut out = FlintInteger::new();
    unsafe {
        ffi::fmpz_euler_phi(out.inner_mut_ptr(), fz.inner_ptr());
    }
    Ok(out.to_string())
}

/// Jacobi symbol \((a | n)\) for odd \(n > 1\) (`fmpz_jacobi`).
pub fn jacobi_symbol(a: &str, n: &str) -> Result<i32, NumberTheoryError> {
    let na = parse_int(a)?;
    let nn = parse_positive(n)?;
    if nn <= 1 || !integer_is_odd(&nn) {
        return Err(NumberTheoryError::Domain {
            msg: "Jacobi denominator must be odd and greater than 1",
        });
    }
    let fa = FlintInteger::from_rug(&na);
    let fn_ = FlintInteger::from_rug(&nn);
    let j = unsafe { ffi::fmpz_jacobi(fa.inner_ptr(), fn_.inner_ptr()) };
    Ok(j as i32)
}

/// Modular \(k\)th root: some \(x\) with \(x^k \equiv a \pmod p\) for prime \(p\).
///
/// Implemented for `k == 2` via `fmpz_sqrtmod`, or when \(\gcd(k, p{-}1)=1\) via exponent inversion.
pub fn nthroot_mod(a: &str, k: u64, p: &str) -> Result<String, NumberTheoryError> {
    if k == 0 {
        return Err(NumberTheoryError::InvalidInput {
            msg: "root degree must be ≥ 1",
        });
    }
    let pm = parse_positive(p)?;
    let fp = FlintInteger::from_rug(&pm);
    if unsafe { ffi::fmpz_is_prime(fp.inner_ptr()) } == 0 {
        return Err(NumberTheoryError::CompositeModulus);
    }

    // Least non-negative residue: rug's `%` truncates toward zero, so a
    // negative `a` would otherwise stay negative.
    let ared = parse_int(a)?.rem_euc(&pm);

    let mut out = FlintInteger::new();

    if k == 2 {
        let fa = FlintInteger::from_rug(&ared);
        let ok = unsafe { ffi::fmpz_sqrtmod(out.inner_mut_ptr(), fa.inner_ptr(), fp.inner_ptr()) };
        if ok == 0 {
            return Err(NumberTheoryError::NoSolution);
        }
        return Ok(out.to_string());
    }

    let ord = pm.clone() - 1;
    let kk = Integer::from(k);
    if kk.clone().gcd(&ord) != 1 {
        return Err(NumberTheoryError::UnsupportedNthRoot);
    }
    let mut inv_e = mod_inverse(kk.clone(), &ord).ok_or(NumberTheoryError::UnsupportedNthRoot)?;
    inv_e %= &ord;
    let fa = FlintInteger::from_rug(&ared);
    let fe = FlintInteger::from_rug(&inv_e);
    unsafe {
        ffi::fmpz_powm(
            out.inner_mut_ptr(),
            fa.inner_ptr(),
            fe.inner_ptr(),
            fp.inner_ptr(),
        );
    }
    Ok(out.to_string())
}

/// Smallest exponent \(e \geq 0\) with \(\mathit{base}^e \equiv \mathit{residue}\pmod{p}\) (`p` prime).
///
/// This uses a deterministic linear sweep over exponents bounded by \(p{-}1\); it is tuned for API
/// parity and moderate primes, not large-field cryptography.
pub fn discrete_log(residue: &str, base: &str, p: &str) -> Result<String, NumberTheoryError> {
    let pm = parse_positive(p)?;
    if pm < 2 {
        return Err(NumberTheoryError::Domain {
            msg: "modulus must be at least 2",
        });
    }
    let fp = FlintInteger::from_rug(&pm);
    if unsafe { ffi::fmpz_is_prime(fp.inner_ptr()) } == 0 {
        return Err(NumberTheoryError::CompositeModulus);
    }

    let ord = pm.clone() - Integer::from(1);
    // Least non-negative residues. rug's `%` truncates toward zero, so
    // `-1 % 5` is `-1`, which no power of the base ever equals: the sweep
    // below used to report NoSolution for `discrete_log(-1, 2, 5)` = 2.
    let b = parse_int(base)?.rem_euc(&pm);
    let r = parse_int(residue)?.rem_euc(&pm);

    if b.is_zero() {
        // 0^0 = 1 and 0^e = 0 for every e >= 1.
        return if r == 1 {
            Ok("0".into())
        } else if r.is_zero() {
            Ok("1".into())
        } else {
            Err(NumberTheoryError::NoSolution)
        };
    }

    let mut cur = Integer::from(1);
    let mut exp = Integer::from(0);
    while exp < ord {
        if cur == r {
            return Ok(exp.to_string());
        }
        cur = (&cur * &b).complete();
        cur %= &pm;
        exp += 1;
    }
    Err(NumberTheoryError::NoSolution)
}

/// Quadratic Dirichlet character: Jacobi symbol \((· | q)\) for odd square-free \(q≥3\).
#[derive(Clone, Debug)]
pub struct QuadraticDirichlet {
    modulus: Integer,
}

impl QuadraticDirichlet {
    pub fn new(conductor: &str) -> Result<Self, NumberTheoryError> {
        let q = parse_positive(conductor)?;
        if q <= 2 || !integer_is_odd(&q) {
            return Err(NumberTheoryError::Domain {
                msg: "quadratic Dirichlet conductor must be odd and ≥ 3",
            });
        }
        let (_sign, fac) = factorint(conductor)?;
        for (_, e) in &fac {
            if *e != 1 {
                return Err(NumberTheoryError::Domain {
                    msg: "conductor must be square-free",
                });
            }
        }
        Ok(QuadraticDirichlet { modulus: q })
    }

    pub fn conductor(&self) -> String {
        self.modulus.to_string()
    }

    /// \(\chi_q(n)\) as `−1`, `0`, or `1`.
    pub fn eval(&self, n: &str) -> Result<i32, NumberTheoryError> {
        jacobi_symbol(n, &self.modulus.to_string())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use rug::ops::Pow;
    use std::collections::HashMap;

    /// A negative residue or base used to stay negative (rug's `%`
    /// truncates), so the sweep never matched: `discrete_log(-1, 2, 5)` said
    /// NoSolution although 2^2 = 4 ≡ -1.
    #[test]
    fn discrete_log_normalises_negative_residues() {
        assert_eq!(discrete_log("-1", "2", "5").unwrap(), "2");
        assert_eq!(discrete_log("4", "-3", "5").unwrap(), "2");
        assert_eq!(discrete_log("-4", "-3", "5").unwrap(), "0");
        assert_eq!(discrete_log("-16", "3", "17").unwrap(), "0");
        assert_eq!(discrete_log("1", "0", "7").unwrap(), "0");
        assert_eq!(discrete_log("0", "0", "7").unwrap(), "1");
        assert_eq!(discrete_log("-7", "7", "7").unwrap(), "1");
        assert!(matches!(
            discrete_log("3", "0", "7"),
            Err(NumberTheoryError::NoSolution)
        ));
    }

    #[test]
    fn nthroot_mod_normalises_negative_input() {
        let x: i64 = nthroot_mod("-1", 2, "5").unwrap().parse().unwrap();
        assert_eq!((x * x) % 5, 4);
        let x: i64 = nthroot_mod("-4", 3, "11").unwrap().parse().unwrap();
        assert_eq!((x * x * x) % 11, 7);
        assert!((0..11).contains(&x));
    }

    #[test]
    fn mersenne_m127_prime() {
        let m = Integer::from(2u32).pow(127_u32) - 1_u32;
        assert!(isprime(&m.to_string()).unwrap());
    }

    #[test]
    fn factorint_zero_is_error() {
        assert!(factorint("0").is_err());
    }

    #[test]
    fn isprime_negative_is_error() {
        assert!(isprime("-5").is_err());
    }

    #[test]
    fn factorint_f5() {
        let n = &(1u128 << 32) - 1;
        let (sign, pairs) = factorint(&n.to_string()).unwrap();
        assert_eq!(sign, 1);
        let m: HashMap<_, _> = pairs.into_iter().collect();
        assert_eq!(m.get("65537").copied(), Some(1));
    }

    /// Audit C6: `fmpz_factor` on three Mersenne primes ran past twenty
    /// seconds under a one-second budget. The ladder stops at the budget.
    #[test]
    fn factorint_bounded_honours_a_wall_budget() {
        use crate::budget::{self, Budget};
        let m = |k: u32| (rug::Integer::from(1) << k) - 1u32;
        let n = m(89) * m(107) * m(127);
        let _guard = budget::enter(Budget::new().with_wall(std::time::Duration::from_millis(500)));
        let t = std::time::Instant::now();
        let r = factorint_bounded(&n.to_string());
        assert!(
            matches!(r, Err(arith::ArithmeticError::Resource(_))),
            "{r:?}"
        );
        assert!(t.elapsed().as_secs_f64() < 5.0, "took {:?}", t.elapsed());
    }

    /// Under a generous budget the ladder factors what `factorint` factors,
    /// including a composite with no small factor ((2⁶¹−1)(2⁸⁹−1), 150 bits —
    /// the full factoriser's range).
    #[test]
    fn factorint_bounded_agrees_with_factorint() {
        use crate::budget::{self, Budget};
        let m = |k: u32| (rug::Integer::from(1) << k) - 1u32;
        let cases = [
            "18446744073709551617".to_string(),
            "340282366920938463463374607431768211457".to_string(),
            "-1000000030000002700000007".to_string(),
            "600851475143".to_string(),
            (m(61) * m(89)).to_string(),
            "10000000000000000000000000000000000000001".to_string(),
        ];
        for n in &cases {
            let want = factorint(n).unwrap();
            let _guard = budget::enter(Budget::new().with_wall(std::time::Duration::from_secs(60)));
            let got = factorint_bounded(n).unwrap();
            assert_eq!(got, want, "{n}");
        }
        // No budget: exactly `factorint`.
        assert_eq!(factorint_bounded("97").unwrap(), factorint("97").unwrap());
        assert!(matches!(
            factorint_bounded("0"),
            Err(arith::ArithmeticError::Input(
                NumberTheoryError::Domain { .. }
            ))
        ));
    }

    #[test]
    fn nextprime_gap() {
        assert_eq!(nextprime("13", true).unwrap(), "17");
    }

    #[test]
    fn totient_twelve() {
        assert_eq!(totient("12").unwrap(), "4");
    }

    #[test]
    fn jacobi_two_fifteen() {
        assert_eq!(jacobi_symbol("2", "15").unwrap(), 1);
    }

    #[test]
    fn sqrt_mod_prime() {
        let x_str = nthroot_mod("144", 2, "401").unwrap();
        let x: u64 = x_str.parse().unwrap();
        assert_eq!((x * x) % 401, 144);
    }

    #[test]
    fn nth_root_via_coprime_exponent() {
        let pm = Integer::from(10007);
        let a = Integer::from(42);
        let k = 5u64;
        let kk = Integer::from(k);
        let ord = pm.clone() - Integer::from(1);
        assert_eq!(kk.clone().gcd(&ord), Integer::from(1));

        let x_str = nthroot_mod(&a.to_string(), k, &pm.to_string()).unwrap();
        let x = Integer::from_str(&x_str).unwrap();
        let chk = x.clone().pow(k as u32) % &pm;
        assert_eq!(chk, a % &pm);
    }

    #[test]
    fn discrete_log_three_mod_seventeen() {
        assert_eq!(discrete_log("13", "3", "17").unwrap(), "4",);
    }

    #[test]
    fn dirichlet_phi_fifteen() {
        let chi = QuadraticDirichlet::new("15").unwrap();
        assert_eq!(chi.eval("14").unwrap(), -1);
        assert_eq!(chi.eval("3").unwrap(), 0);
    }
}