//! Variable-time modular exponentiation for [`U256`].
use super::U256;

impl U256 {
    /// Computes `self^exp mod modulus` in variable time via left-to-right
    /// binary square-and-multiply.
    ///
    /// The loop count depends on `exp.bit_len()` and the conditional multiply
    /// branches on each exponent bit, leaking the exponent through timing.
    /// For cryptographic contexts, use `pow_mod_ct` (enabled via `ct` feature).
    ///
    /// Returns 1 when `exp` is zero (for any base, including zero), matching
    /// the mathematical convention `0^0 = 1`.
    #[inline]
    pub fn pow_mod_vt(&self, exp: &U256, modulus: &U256) -> U256 {
        let bits = exp.bit_len();
        if bits == 0 {
            return U256::ONE;
        }

        let base = self.modulo(modulus);
        let mut result = U256::ONE;

        let mut i = bits;
        while i > 0 {
            i -= 1;
            result = result.square_mod(modulus);
            if exp.bit(i) {
                result = result.mul_mod(&base, modulus);
            }
        }

        result
    }
}

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

    const P: U256 = U256::from_be_limbs([
        0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF,
        0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F,
    ]);

    /// x^0 = 1 for any x.
    #[test]
    fn exp_zero() {
        let base = U256::from_be_limbs([0, 0, 0, 42]);
        assert_eq!(base.pow_mod_vt(&U256::ZERO, &P), U256::ONE);
    }

    /// x^1 = x.
    #[test]
    fn exp_one() {
        let base = U256::from_be_limbs([0, 0, 0, 42]);
        assert_eq!(base.pow_mod_vt(&U256::ONE, &P), base);
    }

    /// 2^10 = 1024.
    #[test]
    fn two_pow_ten() {
        let base = U256::from_be_limbs([0, 0, 0, 2]);
        let exp = U256::from_be_limbs([0, 0, 0, 10]);
        assert_eq!(base.pow_mod_vt(&exp, &P), U256::from_be_limbs([0, 0, 0, 1024]));
    }

    /// Fermat's little theorem: a^(P-1) ≡ 1 (mod P).
    #[test]
    fn fermats_little_theorem() {
        let base = U256::from_be_limbs([0, 0, 0, 7]);
        let p_minus_1 = P - U256::ONE;
        assert_eq!(base.pow_mod_vt(&p_minus_1, &P), U256::ONE);
    }

    /// 0^0 = 1 (mathematical convention).
    #[test]
    fn zero_pow_zero() {
        assert_eq!(U256::ZERO.pow_mod_vt(&U256::ZERO, &P), U256::ONE);
    }
}
