//! Variable-time 512-bit modular reduction for sparse primes.
use super::U512;
use crate::u256::U256;

impl U512 {
    /// Reduces this 512-bit value modulo a sparse prime where the complement
    /// `2^256 − modulus` fits in a single `u64`.
    ///
    /// Exploits the identity `2^256 ≡ comp (mod modulus)` to fold the high
    /// 256 bits into the low 256 bits using just 5 `u64 × u64` multiplies
    /// instead of Knuth Algorithm D's full multi-precision division.
    ///
    /// Round 1 fuses the high-half multiply (`hi × comp`) with the low-half
    /// addition in a single carry-chain pass. Round 2 folds the remaining
    /// overflow (typically ~33 bits for secp256k1). At most 2 conditional
    /// subtractions normalize the result into `[0, modulus)`.
    ///
    /// # Examples
    ///
    /// ```
    /// use cnfy_uint::u256::U256;
    /// use cnfy_uint::u512::U512;
    /// let P = U256::from_be_limbs([
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFFFFFFFFFF,
    ///     0xFFFFFFFFFFFFFFFF, 0xFFFFFFFEFFFFFC2F,
    /// ]);
    ///
    /// let a = U256::from_be_limbs([0, 0, 0, 7]);
    /// let product = a * a;
    /// assert_eq!(product.sparse_reduce_vt(&P, 0x1000003D1), U256::from_be_limbs([0, 0, 0, 49]));
    /// ```
    #[inline]
    pub fn sparse_reduce_vt(&self, modulus: &U256, comp: u64) -> U256 {
        let c = comp as u128;

        // Round 1: lo + hi * comp, fused multiply-accumulate.
        // LE layout: low half = [0..4], high half = [4..8].
        // Each step: self.0[i] + comp * self.0[i+4] + carry.
        // Maximum per step: (2^64-1) + (2^64-1)^2 + (2^64-1) = 2^128-1. Fits in u128.
        let sum = (self.0[0] as u128) + c * (self.0[4] as u128);
        let r0 = sum as u64;
        let carry = sum >> 64;

        let sum = (self.0[1] as u128) + c * (self.0[5] as u128) + carry;
        let r1 = sum as u64;
        let carry = sum >> 64;

        let sum = (self.0[2] as u128) + c * (self.0[6] as u128) + carry;
        let r2 = sum as u64;
        let carry = sum >> 64;

        let sum = (self.0[3] as u128) + c * (self.0[7] as u128) + carry;
        let r3 = sum as u64;
        let overflow = (sum >> 64) as u64;

        // Round 2: fold the overflow. For secp256k1 (comp = 33 bits),
        // overflow is at most ~33 bits and correction is at most ~66 bits.
        if overflow == 0 {
            return U256([r0, r1, r2, r3]).modulo(modulus);
        }

        let correction = (overflow as u128) * c;

        let sum = (r0 as u128) + (correction & 0xFFFF_FFFF_FFFF_FFFF);
        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) as u64;

        let mut result = U256([r0, r1, r2, r3]);

        if (overflow2 != 0) | (result >= *modulus) {
            result = result.overflowing_sub(modulus).0;
        }

        result.modulo(modulus)
    }
}

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

    /// 7^2 = 49, reduced mod P stays 49.
    #[test]
    fn small_product() {
        let a = U256::from_be_limbs([0, 0, 0, 7]);
        let product = a * a;
        assert_eq!(
            product.sparse_reduce_vt(&P, 0x1000003D1),
            U256::from_be_limbs([0, 0, 0, 49])
        );
    }

    /// Zero product reduces to zero.
    #[test]
    fn zero_product() {
        let product = U512::from_be_limbs([0; 8]);
        assert_eq!(product.sparse_reduce_vt(&P, 0x1000003D1), U256::ZERO);
    }

    /// Max product: (2^256-1)^2 mod P matches mul_mod.
    #[test]
    fn max_product() {
        let max = U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]);
        let product = max * max;
        let expected = max.mul_mod(&max, &P);
        assert_eq!(product.sparse_reduce_vt(&P, 0x1000003D1), expected);
    }

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

    /// Sparse reduce matches mul_mod for various products.
    #[test]
    fn matches_mul_mod() {
        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,
                ]),
            ),
        ];
        for (a, b) in &cases {
            let product = *a * *b;
            let expected = a.mul_mod(b, &P);
            let sparse = product.sparse_reduce_vt(&P, 0x1000003D1);
            assert_eq!(sparse, expected, "mismatch for {a:?} * {b:?}");
        }
    }
}
