//! Modular squaring with optimized cross-product and sparse-prime reduction.
use super::U256;

impl U256 {
    /// Computes `self^2 mod modulus` via widening squaring followed by
    /// modular reduction.
    ///
    /// Uses [`square`](U256::square) which requires only
    /// 10 limb multiplications (4 diagonal + 6 cross) instead of the 16
    /// that `Mul<U256>` needs, giving ~37% fewer
    /// multiplies. Reduction uses the same sparse-prime folding technique
    /// as [`mul_mod`](U256::mul_mod) for primes where `2^256 - modulus`
    /// fits in a single `u64`, falling back to Knuth Algorithm D otherwise.
    ///
    /// # 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(&P), U256::from_be_limbs([0, 0, 0, 49]));
    /// ```
    #[inline]
    pub fn square_mod(&self, modulus: &U256) -> U256 {
        // Fast path: fused square-reduce for sparse primes.
        // Avoids U512 intermediate by inlining squaring + sparse_reduce.
        if let Some(comp) = modulus.has_u64_complement() {
            return self.square_mod_sparse(modulus, comp);
        }

        // Generic path: widening square then Knuth Algorithm D
        let product = self.square();
        product.div_u256(modulus)
    }
}

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

    /// 0^2 mod P = 0.
    #[test]
    fn square_zero() {
        assert_eq!(U256::ZERO.square_mod(&P), U256::ZERO);
    }

    /// 1^2 mod P = 1.
    #[test]
    fn square_one() {
        assert_eq!(U256::ONE.square_mod(&P), U256::ONE);
    }

    /// 7^2 mod P = 49.
    #[test]
    fn small_square() {
        let a = U256::from_be_limbs([0, 0, 0, 7]);
        assert_eq!(a.square_mod(&P), U256::from_be_limbs([0, 0, 0, 49]));
    }

    /// square_mod matches mul_mod(self, self) for various inputs.
    #[test]
    fn matches_mul_mod() {
        let cases = [
            U256::from_be_limbs([0, 0, 0, 42]),
            U256::from_be_limbs([0, 0, 0x12345678, 0x9ABCDEF0]),
            U256::from_be_limbs([0x1111, 0x2222, 0x3333, 0x4444]),
            U256::from_be_limbs([
                0x123456789ABCDEF0,
                0xFEDCBA9876543210,
                0x1111222233334444,
                0x5555666677778888,
            ]),
            U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]),
            U256::from_be_limbs([
                0xFFFFFFFFFFFFFFFF,
                0xFFFFFFFFFFFFFFFF,
                0xFFFFFFFFFFFFFFFF,
                0xFFFFFFFEFFFFFC2E,
            ]),
        ];
        for a in &cases {
            assert_eq!(a.square_mod(&P), a.mul_mod(a, &P), "mismatch for {a:?}",);
        }
    }

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

    /// Sparse path matches generic path for the max product.
    #[test]
    fn sparse_matches_generic() {
        let a = U256::from_be_limbs([u64::MAX, u64::MAX, u64::MAX, u64::MAX]);
        let product = a.square();
        let generic = product.div_u256(&P);
        let sparse = product.sparse_reduce(&P, 0x1000003D1);
        assert_eq!(sparse, generic);
    }

    /// Unreduced input is reduced correctly.
    #[test]
    fn unreduced_input() {
        let max = U256::from_be_bytes([0xFF; 32]);
        assert_eq!(max.square_mod(&P), max.mul_mod(&max, &P));
    }
}
