//! Constant-time fused multiply-and-reduce for sparse primes.
use super::U256;

impl U256 {
    /// Computes `(self * other) mod modulus` for sparse primes in constant time.
    ///
    /// Same algorithm as `mul_mod_sparse_vt` but with all data-dependent branches
    /// removed. Round 2 always executes (when overflow is zero, correction is zero
    /// so the additions are no-ops). The final reduction uses branchless
    /// conditional subtractions via `ct_select` instead of the `modulo` while loop.
    #[inline]
    pub fn mul_mod_sparse_ct(&self, other: &U256, modulus: &U256, comp: u64) -> U256 {
        let a = self.0;
        let b = other.0;

        // === Phase 1: Schoolbook 4x4 multiply (identical to VT) ===

        let p = (a[0] as u128) * (b[0] as u128);
        let r0 = p as u64;
        let carry = p >> 64;

        let p = (a[0] as u128) * (b[1] as u128) + carry;
        let r1 = p as u64;
        let carry = p >> 64;

        let p = (a[0] as u128) * (b[2] as u128) + carry;
        let r2 = p as u64;
        let carry = p >> 64;

        let p = (a[0] as u128) * (b[3] as u128) + carry;
        let r3 = p as u64;
        let r4 = (p >> 64) as u64;

        let p = (a[1] as u128) * (b[0] as u128) + (r1 as u128);
        let r1 = p as u64;
        let carry = p >> 64;

        let p = (a[1] as u128) * (b[1] as u128) + (r2 as u128) + carry;
        let r2 = p as u64;
        let carry = p >> 64;

        let p = (a[1] as u128) * (b[2] as u128) + (r3 as u128) + carry;
        let r3 = p as u64;
        let carry = p >> 64;

        let p = (a[1] as u128) * (b[3] as u128) + (r4 as u128) + carry;
        let r4 = p as u64;
        let r5 = (p >> 64) as u64;

        let p = (a[2] as u128) * (b[0] as u128) + (r2 as u128);
        let r2 = p as u64;
        let carry = p >> 64;

        let p = (a[2] as u128) * (b[1] as u128) + (r3 as u128) + carry;
        let r3 = p as u64;
        let carry = p >> 64;

        let p = (a[2] as u128) * (b[2] as u128) + (r4 as u128) + carry;
        let r4 = p as u64;
        let carry = p >> 64;

        let p = (a[2] as u128) * (b[3] as u128) + (r5 as u128) + carry;
        let r5 = p as u64;
        let r6 = (p >> 64) as u64;

        let p = (a[3] as u128) * (b[0] as u128) + (r3 as u128);
        let r3 = p as u64;
        let carry = p >> 64;

        let p = (a[3] as u128) * (b[1] as u128) + (r4 as u128) + carry;
        let r4 = p as u64;
        let carry = p >> 64;

        let p = (a[3] as u128) * (b[2] as u128) + (r5 as u128) + carry;
        let r5 = p as u64;
        let carry = p >> 64;

        let p = (a[3] as u128) * (b[3] as u128) + (r6 as u128) + carry;
        let r6 = p as u64;
        let r7 = (p >> 64) as u64;

        // === Phase 2: Sparse reduce (always both rounds) ===

        let c = comp as u128;

        // Round 1: fused multiply-accumulate of hi * comp + lo
        let sum = (r0 as u128) + c * (r4 as u128);
        let r0 = sum as u64;
        let carry = sum >> 64;

        let sum = (r1 as u128) + c * (r5 as u128) + carry;
        let r1 = sum as u64;
        let carry = sum >> 64;

        let sum = (r2 as u128) + c * (r6 as u128) + carry;
        let r2 = sum as u64;
        let carry = sum >> 64;

        let sum = (r3 as u128) + c * (r7 as u128) + carry;
        let r3 = sum as u64;
        let overflow = (sum >> 64) as u64;

        // Round 2: fold remaining overflow (always executed; no-op when overflow=0)
        let correction = (overflow as u128) * c;

        let sum = (r0 as u128) + (correction as u64 as u128);
        let r0 = sum as u64;
        let carry = sum >> 64;

        let sum = (r1 as u128) + (correction >> 64) + carry;
        let r1 = sum as u64;
        let carry = sum >> 64;

        let sum = (r2 as u128) + carry;
        let r2 = sum as u64;
        let carry = sum >> 64;

        let sum = (r3 as u128) + carry;
        let r3 = sum as u64;
        let overflow2 = (sum >> 64) != 0;

        // === Phase 3: Branchless final reduction ===
        let result = U256([r0, r1, r2, r3]);

        // First conditional subtraction (accounts for overflow2 bit)
        let (r1_sub, borrow1) = result.overflowing_sub(modulus);
        let after1 = r1_sub.ct_select(&result, overflow2 | !borrow1);

        // Second conditional subtraction (ensures result < modulus)
        let (r2_sub, borrow2) = after1.overflowing_sub(modulus);
        r2_sub.ct_select(&after1, !borrow2)
    }
}

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

    /// Small product: 7 * 6 = 42.
    #[test]
    fn small_product() {
        let a = U256::from_be_limbs([0, 0, 0, 7]);
        let b = U256::from_be_limbs([0, 0, 0, 6]);
        assert_eq!(a.mul_mod_sparse_ct(&b, &P, COMP), U256::from_be_limbs([0, 0, 0, 42]));
    }

    /// Zero times anything is zero.
    #[test]
    fn zero_times_any() {
        let a = U256::from_be_limbs([0x1234, 0x5678, 0x9ABC, 0xDEF0]);
        assert_eq!(U256::ZERO.mul_mod_sparse_ct(&a, &P, COMP), U256::ZERO);
    }

    /// (P - 1)^2 mod P = 1.
    #[test]
    fn p_minus_one_squared() {
        let p_minus_1 = U256::from_be_limbs([
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFFFFFFFFFF,
            0xFFFFFFFEFFFFFC2E,
        ]);
        assert_eq!(p_minus_1.mul_mod_sparse_ct(&p_minus_1, &P, COMP), U256::ONE);
    }

    /// Multiplicative identity.
    #[test]
    fn multiplicative_identity() {
        let a = U256::from_be_limbs([0x1234, 0x5678, 0x9ABC, 0xDEF0]);
        assert_eq!(U256::ONE.mul_mod_sparse_ct(&a, &P, COMP), a);
    }

    /// CT matches VT for diverse inputs.
    #[test]
    fn matches_vt() {
        let cases = [
            (U256::from_be_limbs([0, 0, 0, 1]), U256::from_be_limbs([0, 0, 0, 1])),
            (
                U256::from_be_limbs([
                    0x123456789ABCDEF0, 0xFEDCBA9876543210,
                    0x1111222233334444, 0x5555666677778888,
                ]),
                U256::from_be_limbs([
                    0xAAAABBBBCCCCDDDD, 0xEEEEFFFF00001111,
                    0x2222333344445555, 0x6666777788889999,
                ]),
            ),
            (
                U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, 0xFFFFFFFEFFFFFC2E]),
                U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, 0xFFFFFFFEFFFFFC2E]),
            ),
            (
                U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]),
                U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]),
            ),
            (
                U256::from_be_limbs([
                    0x79BE667EF9DCBBAC, 0x55A06295CE870B07,
                    0x029BFCDB2DCE28D9, 0x59F2815B16F81798,
                ]),
                U256::from_be_limbs([
                    0x483ADA7726A3C465, 0x5DA4FBFC0E1108A8,
                    0xFD17B448A6855419, 0x9C47D08FFB10D4B8,
                ]),
            ),
        ];
        for (a, b) in &cases {
            assert_eq!(
                a.mul_mod_sparse_ct(b, &P, COMP),
                a.mul_mod_sparse_vt(b, &P, COMP),
                "CT/VT mismatch for {a:?} * {b:?}"
            );
        }
    }
}
