//! Constant-time modular reduction via three branchless conditional subtractions.
use super::U256;

impl U256 {
    /// Reduces `self` modulo `modulus` in constant time using three branchless
    /// conditional subtractions.
    ///
    /// Handles inputs up to `3 * modulus - 1`, which covers all values produced
    /// by sparse-prime reduction paths. Each pass performs an unconditional
    /// subtraction followed by a `ct_select` to keep the reduced value only if
    /// no underflow occurred, ensuring execution time is independent of the
    /// input magnitude.
    ///
    /// # Panics
    ///
    /// Produces incorrect results if `self >= 4 * modulus`.
    #[inline]
    pub fn modulo_ct(&self, modulus: &U256) -> U256 {
        let (r1, borrow1) = self.overflowing_sub(modulus);
        let val1 = r1.ct_select(self, !borrow1);

        let (r2, borrow2) = val1.overflowing_sub(modulus);
        let val2 = r2.ct_select(&val1, !borrow2);

        let (r3, borrow3) = val2.overflowing_sub(modulus);
        r3.ct_select(&val2, !borrow3)
    }
}

#[cfg(test)]
mod ai_tests {
    use super::*;
    const P: U256 = U256::from_be_limbs([
        0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF,
        0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F,
    ]);

    /// Zero mod P is zero.
    #[test]
    fn zero_mod_p() {
        assert_eq!(U256::ZERO.modulo_ct(&P), U256::ZERO);
    }

    /// Value less than modulus is unchanged.
    #[test]
    fn below_modulus() {
        assert_eq!(U256::ONE.modulo_ct(&P), U256::ONE);
    }

    /// Modulus mod itself is zero.
    #[test]
    fn modulus_mod_self() {
        assert_eq!(P.modulo_ct(&P), U256::ZERO);
    }

    /// (2^256 - 1) mod P = 2^32 + 976 (value is P + complement).
    #[test]
    fn max_mod_p() {
        let max = U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]);
        assert_eq!(max.modulo_ct(&P), U256::from_be_limbs([0, 0, 0, 0x1000003D0]));
    }

    /// 2M reduces to zero for a small modulus where 2M fits in 256 bits.
    #[test]
    fn double_modulus_small() {
        let m = U256::from_be_limbs([0, 0, 0, 1000]);
        let two_m = U256::from_be_limbs([0, 0, 0, 2000]);
        assert_eq!(two_m.modulo_ct(&m), U256::ZERO);
    }

    /// 2M + 1 reduces to 1 for a small modulus.
    #[test]
    fn double_modulus_plus_one_small() {
        let m = U256::from_be_limbs([0, 0, 0, 1000]);
        let v = U256::from_be_limbs([0, 0, 0, 2001]);
        assert_eq!(v.modulo_ct(&m), U256::ONE);
    }

    /// 3M - 1 reduces to M - 1 for a small modulus.
    #[test]
    fn triple_modulus_minus_one_small() {
        let m = U256::from_be_limbs([0, 0, 0, 1000]);
        let v = U256::from_be_limbs([0, 0, 0, 2999]);
        let expected = U256::from_be_limbs([0, 0, 0, 999]);
        assert_eq!(v.modulo_ct(&m), expected);
    }

    /// P + small values (typical sparse-reduction output) reduce correctly.
    #[test]
    fn p_plus_small() {
        let (p_plus_1, _) = P.overflowing_add(&U256::ONE);
        assert_eq!(p_plus_1.modulo_ct(&P), U256::ONE);

        let comp = U256::from_be_limbs([0, 0, 0, 0x1000003D0]);
        let (p_plus_comp, _) = P.overflowing_add(&comp);
        assert_eq!(p_plus_comp.modulo_ct(&P), comp);
    }

    /// CT matches VT for reduced inputs.
    #[test]
    fn matches_vt() {
        let cases = [
            U256::ZERO,
            U256::ONE,
            U256::from_be_limbs([0, 0, 0, 42]),
            U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]),
            P,
        ];
        for v in &cases {
            assert_eq!(v.modulo_ct(&P), v.modulo_vt(&P), "mismatch for {v:?}");
        }
    }
}
