833: Discrete Logarithm — Baby-Step Giant-Step

/// Discrete Logarithm: Baby-Step Giant-Step (BSGS).
///
/// Solves a^x ≡ b (mod p) for prime p in O(√p) time and space.
/// Write x = i·m - j, precompute baby steps, check giant steps.

use std::collections::HashMap;

fn mulmod(a: u64, b: u64, m: u64) -> u64 {
    (a as u128 * b as u128 % m as u128) as u64
}

fn pow_mod(mut base: u64, mut exp: u64, m: u64) -> u64 {
    let mut result = 1u64;
    base %= m;
    while exp > 0 {
        if exp & 1 == 1 { result = mulmod(result, base, m); }
        base = mulmod(base, base, m);
        exp >>= 1;
    }
    result
}

/// BSGS: find smallest x ≥ 0 such that a^x ≡ b (mod p).
/// p must be prime (or at least gcd(a, p) = 1).
/// Returns None if no solution exists.
pub fn bsgs(a: u64, b: u64, p: u64) -> Option<u64> {
    let m = (p as f64).sqrt().ceil() as u64;

    // Baby steps: store a^j mod p → j for j = 0..m
    let mut table: HashMap<u64, u64> = HashMap::with_capacity(m as usize + 1);
    let mut aj = 1u64;
    for j in 0..=m {
        // Only store first occurrence (smallest j wins)
        table.entry(aj).or_insert(j);
        aj = mulmod(aj, a, p);
    }

    // Giant steps: check b * (a^m)^i mod p for i = 1..=m
    let am = pow_mod(a, m, p);
    let mut giant = b;
    for i in 1..=m {
        giant = mulmod(giant, am, p);
        if let Some(&j) = table.get(&giant) {
            let x = i * m - j;
            // Verify: x must be non-negative (it always is here since i*m ≥ j)
            return Some(x);
        }
    }
    None
}

/// Brute-force for verification and small p.
fn discrete_log_brute(a: u64, b: u64, p: u64) -> Option<u64> {
    let mut cur = 1u64;
    for x in 0..p {
        if cur == b { return Some(x); }
        cur = mulmod(cur, a, p);
    }
    None
}

fn main() {
    let cases = [
        (2u64, 22, 29),  // 2^x ≡ 22 (mod 29)
        (3u64, 1, 7),    // 3^x ≡ 1 (mod 7): x=6 (order of 3)
        (5u64, 3, 23),   // 5^x ≡ 3 (mod 23)
        (2u64, 3, 7),    // 2^x ≡ 3 (mod 7)
        (7u64, 2, 11),   // 7^x ≡ 2 (mod 11)
    ];

    for (a, b, p) in cases {
        let result = bsgs(a, b, p);
        let brute = discrete_log_brute(a, b, p);
        let verify = result.map(|x| pow_mod(a, x, p) == b).unwrap_or(true);
        println!("{a}^x ≡ {b} (mod {p}): bsgs={result:?}, brute={brute:?}, verified={verify}");
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn check(a: u64, b: u64, p: u64) {
        let bsgs_result = bsgs(a, b, p);
        let brute_result = discrete_log_brute(a, b, p);
        // Both should agree on existence
        assert_eq!(bsgs_result.is_some(), brute_result.is_some(),
            "existence mismatch: {a}^x≡{b} mod {p}");
        // If solution exists, verify it
        if let Some(x) = bsgs_result {
            assert_eq!(pow_mod(a, x, p), b,
                "wrong answer: {a}^{x} mod {p} ≠ {b}");
        }
    }

    #[test]
    fn test_basic_cases() {
        check(2, 22, 29);
        check(3, 1, 7);
        check(5, 3, 23);
        check(2, 3, 7);
    }

    #[test]
    fn test_identity() {
        // a^0 = 1
        let result = bsgs(5, 1, 23);
        // Could be 0 or 22 (the order)
        assert!(result.is_some());
        let x = result.unwrap();
        assert_eq!(pow_mod(5, x, 23), 1);
    }

    #[test]
    fn test_large_prime() {
        // 2^x ≡ 100 (mod 1009)
        check(2, 100, 1009);
    }

    #[test]
    fn test_matches_brute_all_small() {
        let p = 29u64; // Small prime
        for b in 1..p {
            check(2, b, p);
        }
    }

    #[test]
    fn test_no_solution() {
        // If a has order d < p-1, some b values have no log
        // 2 is not a primitive root mod 9 (which is composite), but for primes
        // we can still have no solution if b is not in the subgroup generated by a
        // (When a is a primitive root, all b have a log.)
        // Here we just verify our function doesn't panic or lie.
        let result = bsgs(3, 5, 7); // check
        let brute = discrete_log_brute(3, 5, 7);
        assert_eq!(result.is_some(), brute.is_some());
    }
}