//! Constant-time modular addition of two [`U256`] values.
use super::U256;

impl U256 {
    /// Computes `(self + other) mod modulus` in constant time.
    ///
    /// Both operands must be in `[0, modulus)` (reduced). For reduced inputs
    /// the sum is at most `2 * modulus - 2`, so a single conditional
    /// subtraction always suffices. The subtraction is always executed and
    /// the result selected via branchless bitmask (`ct_select`), ensuring
    /// execution time is independent of the operand values.
    ///
    /// # Examples
    ///
    /// ```
    /// use cnfy_uint::u256::U256;
    /// let P = U256::from_be_limbs([
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF,
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F,
    /// ]);
    ///
    /// let a = U256::from_be_limbs([0, 0, 0, 10]);
    /// let b = U256::from_be_limbs([0, 0, 0, 20]);
    /// assert_eq!(a.add_mod_ct(&b, &P), U256::from_be_limbs([0, 0, 0, 30]));
    /// ```
    #[inline]
    pub fn add_mod_ct(&self, other: &U256, modulus: &U256) -> U256 {
        let (sum, carry) = self.overflowing_add(other);
        let (reduced, borrow) = sum.overflowing_sub(modulus);

        // Subtraction is valid when carry absorbs the borrow (overflow case)
        // or when sum >= modulus (no borrow). One subtraction always suffices
        // for reduced inputs (both < modulus), since their sum < 2 * modulus.
        reduced.ct_select(&sum, carry | !borrow)
    }
}

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

    /// Adding zero to zero yields zero.
    #[test]
    fn zero_plus_zero() {
        assert_eq!(U256::ZERO.add_mod_ct(&U256::ZERO, &P), U256::ZERO);
    }

    /// Adding zero is the identity operation.
    #[test]
    fn additive_identity() {
        let a = U256::from_be_limbs([0, 0, 0, 42]);
        assert_eq!(a.add_mod_ct(&U256::ZERO, &P), a);
        assert_eq!(U256::ZERO.add_mod_ct(&a, &P), a);
    }

    /// Small values that don't require reduction.
    #[test]
    fn small_no_reduction() {
        let a = U256::from_be_limbs([
            0x0202020202020202,
            0x0202020202020202,
            0x0202020202020202,
            0x0202020202020202,
        ]);
        let expected = U256::from_be_limbs([
            0x0404040404040404,
            0x0404040404040404,
            0x0404040404040404,
            0x0404040404040404,
        ]);
        assert_eq!(a.add_mod_ct(&a, &P), expected);
    }

    /// Adding one to P-1 wraps to zero (exact modulus hit).
    #[test]
    fn p_minus_one_plus_one() {
        let p_minus_1 = U256::from_be_limbs([
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFEFFFFFC2E,
        ]);
        assert_eq!(p_minus_1.add_mod_ct(&U256::ONE, &P), U256::ZERO);
    }

    /// Adding two to P-1 wraps to one.
    #[test]
    fn p_minus_one_plus_two() {
        let p_minus_1 = U256::from_be_limbs([
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFEFFFFFC2E,
        ]);
        let two = U256::from_be_limbs([0, 0, 0, 2]);
        assert_eq!(p_minus_1.add_mod_ct(&two, &P), U256::ONE);
    }

    /// (P-1) + (P-1) mod P = P-2. Exercises 257-bit overflow path.
    #[test]
    fn p_minus_one_doubled() {
        let p_minus_1 = U256::from_be_limbs([
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFEFFFFFC2E,
        ]);
        let p_minus_2 = U256::from_be_limbs([
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFEFFFFFC2D,
        ]);
        assert_eq!(p_minus_1.add_mod_ct(&p_minus_1, &P), p_minus_2);
    }

    /// Commutativity: a + b = b + a.
    #[test]
    fn commutative() {
        let a = U256::from_be_limbs([0, 0x1234, 0, 0x5678]);
        let b = U256::from_be_limbs([0, 0xABCD, 0, 0xEF01]);
        assert_eq!(a.add_mod_ct(&b, &P), b.add_mod_ct(&a, &P));
    }

    /// Two large reduced values whose sum overflows 256 bits.
    #[test]
    fn large_reduced_overflow() {
        let p_minus_1 = U256::from_be_limbs([
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFEFFFFFC2E,
        ]);
        let p_minus_2 = U256::from_be_limbs([
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFEFFFFFC2D,
        ]);
        let expected = U256::from_be_limbs([
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFEFFFFFC2C,
        ]);
        assert_eq!(p_minus_1.add_mod_ct(&p_minus_2, &P), expected);
    }

    /// CT version matches VT version for all reduced inputs.
    #[test]
    fn matches_vt() {
        let cases = [
            (U256::ZERO, U256::ZERO),
            (U256::ONE, U256::ONE),
            (
                U256::from_be_limbs([0, 0x1234, 0, 0x5678]),
                U256::from_be_limbs([0, 0xABCD, 0, 0xEF01]),
            ),
            (
                U256::from_be_limbs([
                    0xFFFFFFFFFFFFFFFF,
                    0xFFFFFFFFFFFFFFFF,
                    0xFFFFFFFFFFFFFFFF,
                    0xFFFFFFFEFFFFFC2E,
                ]),
                U256::from_be_limbs([
                    0xFFFFFFFFFFFFFFFF,
                    0xFFFFFFFFFFFFFFFF,
                    0xFFFFFFFFFFFFFFFF,
                    0xFFFFFFFEFFFFFC2E,
                ]),
            ),
            (
                U256::from_be_limbs([
                    0x79BE667EF9DCBBAC,
                    0x55A06295CE870B07,
                    0x029BFCDB2DCE28D9,
                    0x59F2815B16F81798,
                ]),
                U256::from_be_limbs([
                    0x483ADA7726A3C465,
                    0x5DA4FBFC0E1108A8,
                    0xFD17B448A6855419,
                    0x9C47D08FFB10D4B8,
                ]),
            ),
        ];
        for (a, b) in &cases {
            assert_eq!(
                a.add_mod_ct(b, &P),
                a.add_mod_vt(b, &P),
                "CT/VT mismatch for {a:?} + {b:?}"
            );
        }
    }
}
