//! Wrapping 256-bit multiplication for coefficient updates.

use super::GcdState;
use crate::u256::U256;

impl GcdState {
    /// Wrapping 256-bit multiplication: returns the lower 256 bits of `a × self.t1`.
    ///
    /// Used by [`full_step`](Self::full_step) (cold path, called 1–2 times per
    /// inversion) for the coefficient update `t0 - q × t1` where `q` is a
    /// full U256. Decomposes `self.t1` into limbs, computes four scalar × U256
    /// products, then accumulates only the limbs that contribute to the
    /// lower 256 bits.
    pub(super) fn mul_u256(&self, a: &U256) -> U256 {
        let a_limbs = a.to_le_limbs();
        let b_limbs = self.t1.to_le_limbs();

        let p0 = (*a * b_limbs[0]).to_le_limbs(); // shift 0
        let p1 = (*a * b_limbs[1]).to_le_limbs(); // shift 64
        let p2 = (*a * b_limbs[2]).to_le_limbs(); // shift 128
        // shift 192: only LSB of a contributes to result[3]
        let p3_lsb = ((a_limbs[0] as u128) * (b_limbs[3] as u128)) as u64;

        let r0 = p0[0];

        let c = (p0[1] as u128) + (p1[0] as u128);
        let r1 = c as u64;
        let carry = (c >> 64) as u64;

        let c = (p0[2] as u128) + (p1[1] as u128) + (p2[0] as u128) + (carry as u128);
        let r2 = c as u64;
        let carry = (c >> 64) as u64;

        // Top limb: wrapping (carry out discarded)
        let r3 = p0[3]
            .wrapping_add(p1[2])
            .wrapping_add(p2[1])
            .wrapping_add(p3_lsb)
            .wrapping_add(carry);

        U256([r0, r1, r2, r3])
    }
}

#[cfg(test)]
mod ai_tests {
    use crate::u256::U256;

    use super::*;

    /// Helper: build a GcdState with a specific t1 value.
    fn state_with_t1(t1: U256) -> GcdState {
        GcdState {
            modulus: U256::ZERO,
            r0: U256::ZERO,
            r1: U256::ZERO,
            t0: U256::ZERO,
            t1,
            even: true,
        }
    }

    /// Small values multiply correctly.
    #[test]
    fn small_values() {
        let state = state_with_t1(U256::from_be_limbs([0, 0, 0, 13]));
        let a = U256::from_be_limbs([0, 0, 0, 7]);
        assert_eq!(state.mul_u256(&a), U256::from_be_limbs([0, 0, 0, 91]));
    }

    /// Overflow wraps to lower 256 bits.
    #[test]
    fn overflow_wraps() {
        // MAX * 2 = 2^257 - 2 ≡ 2^256 - 2 mod 2^256 = MAX - 1
        let state = state_with_t1(U256::from_be_limbs([0, 0, 0, 2]));
        let result = state.mul_u256(&U256::MAX);
        let expected = U256::MAX.overflowing_sub(&U256::ONE).0;
        assert_eq!(result, expected);
    }

    /// Multiplying by one is identity.
    #[test]
    fn multiply_by_one() {
        let t1 = U256::from_be_limbs([0x1234, 0x5678, 0x9ABC, 0xDEF0]);
        let state = state_with_t1(t1);
        assert_eq!(state.mul_u256(&U256::ONE), t1);
    }

    /// Multiplying by zero is zero.
    #[test]
    fn multiply_by_zero() {
        let state = state_with_t1(U256::from_be_limbs([0x1234, 0x5678, 0x9ABC, 0xDEF0]));
        assert_eq!(state.mul_u256(&U256::ZERO), U256::ZERO);
    }

    /// Cross-limb multiplication produces correct lower 256 bits.
    #[test]
    fn cross_limb() {
        // 2^128 * 2^64 = 2^192
        let state = state_with_t1(U256::from_be_limbs([0, 0, 1, 0])); // t1 = 2^64
        let a = U256::from_be_limbs([0, 1, 0, 0]); // a = 2^128
        assert_eq!(state.mul_u256(&a), U256::from_be_limbs([1, 0, 0, 0]));
    }
}
