//! Fused square-and-reduce for sparse primes, avoiding the U512 intermediate.
//!
//! Dispatches to `square_mod_sparse_ct` (constant-time) when the `ct-field`
//! feature is enabled, or `square_mod_sparse_vt` (variable-time) by default.
use super::U256;

impl U256 {
    /// Computes `self^2 mod modulus` for sparse primes in a single fused pass.
    ///
    /// `comp` is the sparse prime complement: `2^256 - modulus`, which must
    /// fit in a single `u64`.
    ///
    /// With the `ct-field` feature enabled, this is constant-time (branchless
    /// reduction). Without the feature, uses early-return optimizations.
    ///
    /// # 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, 7]);
    /// assert_eq!(a.square_mod_sparse(&P, 0x1000003D1), U256::from_be_limbs([0, 0, 0, 49]));
    /// ```
    #[inline]
    pub fn square_mod_sparse(&self, modulus: &U256, comp: u64) -> U256 {
        #[cfg(feature = "ct-field")]
        { self.square_mod_sparse_ct(modulus, comp) }
        #[cfg(not(feature = "ct-field"))]
        { self.square_mod_sparse_vt(modulus, comp) }
    }
}

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

    /// Dispatcher matches VT variant.
    #[test]
    fn matches_vt() {
        let cases = [
            U256::from_be_limbs([0, 0, 0, 7]),
            U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, 0xFFFFFFFEFFFFFC2E]),
        ];
        for a in &cases {
            assert_eq!(
                a.square_mod_sparse(&P, COMP),
                a.square_mod_sparse_vt(&P, COMP),
            );
        }
    }

    /// Commutativity with mul_mod_sparse.
    #[test]
    fn matches_mul_mod_sparse() {
        let a = U256::from_be_limbs([
            0xAAAABBBBCCCCDDDD, 0xEEEEFFFF00001111,
            0x2222333344445555, 0x6666777788889999,
        ]);
        assert_eq!(
            a.square_mod_sparse(&P, COMP),
            a.mul_mod_sparse(&a, &P, COMP)
        );
    }
}
