arcium-primitives 0.8.5

use std::num::Wrapping;

use itertools::Itertools;

/// Whether this CPU carries out a carry-less multiply itself.
///
/// Checked at run time rather than through `cfg(target_feature)` alone, which resolves against
/// whatever the build opted into: baseline x86-64 does not include `pclmulqdq`, so gating on it
/// compiles the intrinsics out silently and leaves every Gf2_128 multiply on the quadratic
/// fallback below.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
static HAS_CLMUL: std::sync::LazyLock<bool> = std::sync::LazyLock::new(|| {
    std::arch::is_x86_feature_detected!("pclmulqdq") && std::arch::is_x86_feature_detected!("sse2")
});

/// Whether PCLMULQDQ is available.
///
/// A build that already targets the features MUST NOT pay for a run-time check: the `cfg!` is a
/// compile-time constant, so the branch and the `LazyLock` read fold away and the intrinsic still
/// inlines. Reading the `LazyLock` on every call costs about 3% on the binary triples path.
///
/// `inline(always)` earns its place only in unoptimized builds. At `-O` this inlines either way
/// and the symbol is never emitted; without the attribute a debug build makes a real call per
/// multiply.
#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[inline(always)]
fn has_clmul() -> bool {
    if cfg!(all(target_feature = "pclmulqdq", target_feature = "sse2")) {
        true
    } else {
        *HAS_CLMUL
    }
}

/// Whether this CPU carries PMULL, which lives behind Armv8's optional `aes` extension.
///
/// `neon` is baseline on aarch64 and `aes` is not, so a default build cannot assume PMULL is
/// there. `aarch64-apple-darwin` is the exception: it carries `aes` in its baseline, so there
/// [`has_pmull`] folds away instead.
///
/// Both features are detected, not just `aes`, because both are what the NEON functions declare
/// through `#[target_feature]` and therefore what calling them requires.
#[cfg(target_arch = "aarch64")]
static HAS_PMULL: std::sync::LazyLock<bool> = std::sync::LazyLock::new(|| {
    std::arch::is_aarch64_feature_detected!("neon")
        && std::arch::is_aarch64_feature_detected!("aes")
});

/// Whether PMULL is available. See [`has_clmul`] for why it is shaped this way.
#[cfg(target_arch = "aarch64")]
#[inline(always)]
fn has_pmull() -> bool {
    if cfg!(all(target_feature = "neon", target_feature = "aes")) {
        true
    } else {
        *HAS_PMULL
    }
}

// Carry-less multiplication of two n-limb numbers.
// Output has 2.n limbs.
pub fn carry_less_mul<const LIMBS: usize>(
    a: [u64; LIMBS],
    b: [u64; LIMBS],
) -> ([u64; LIMBS], [u64; LIMBS]) {
    let mut res = [[0u64; LIMBS]; 2];
    let slice = res.as_flattened_mut();

    (0..a.len())
        .cartesian_product(0..b.len())
        .for_each(|(i, j)| clfma_u64(&mut slice[i + j..i + j + 2], a[i], b[j]));

    (res[0], res[1])
}

pub fn carry_less_mul_1limb(a: [u64; 1], b: [u64; 1]) -> ([u64; 1], [u64; 1]) {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    if has_clmul() {
        // SAFETY: the features the function needs are the ones just checked for.
        return unsafe { clmul_u64_sse2(a, b) };
    }

    #[cfg(target_arch = "aarch64")]
    if has_pmull() {
        // SAFETY: the features the function needs are the ones just checked for.
        return unsafe { clmul_u64_neon(a, b) };
    }

    #[allow(unreachable_code)]
    carry_less_mul(a, b)
}

pub fn carry_less_mul_2limbs(a: [u64; 2], b: [u64; 2]) -> ([u64; 2], [u64; 2]) {
    #[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
    if has_clmul() {
        // SAFETY: the features the function needs are the ones just checked for.
        return unsafe { clmul_2_u64_sse2(a, b) };
    }

    #[cfg(target_arch = "aarch64")]
    if has_pmull() {
        // SAFETY: the features the function needs are the ones just checked for.
        return unsafe { clmul_2_u64_neon(a, b) };
    }

    #[allow(unreachable_code)]
    carry_less_mul(a, b)
}

fn clfma_u64(res: &mut [u64], a: u64, b: u64) {
    let bit_k: u64 = (-Wrapping((a) & 0x01)).0;
    res[0] ^= b & bit_k;

    for k in 1..64 {
        let bit_k: u64 = (-Wrapping((a >> k) & 0x01)).0;
        let tmp = b & bit_k;
        res[0] ^= tmp << k;
        res[1] ^= tmp >> (64 - k);
    }
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "pclmulqdq,sse2")]
unsafe fn clmul_2_u64_sse2(x: [u64; 2], y: [u64; 2]) -> ([u64; 2], [u64; 2]) {
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;
    use std::mem::MaybeUninit;

    unsafe {
        let a_b = _mm_loadu_si128(x.as_ptr() as *const _);
        let c_d = _mm_loadu_si128(y.as_ptr() as *const _);

        let prod_a_c = _mm_clmulepi64_si128::<0>(a_b, c_d);
        let prod_b_d = _mm_clmulepi64_si128::<17>(a_b, c_d);

        let s_ab = _mm_xor_si128(_mm_unpacklo_epi64(a_b, a_b), a_b);
        let s_cd = _mm_xor_si128(_mm_unpacklo_epi64(c_d, c_d), c_d);
        let prod_ab_cd = _mm_clmulepi64_si128::<17>(s_ab, s_cd);

        let prod_ab_cd = _mm_xor_si128(prod_ab_cd, _mm_xor_si128(prod_a_c, prod_b_d));

        let res_l = _mm_xor_si128(prod_a_c, _mm_slli_si128::<8>(prod_ab_cd));
        let res_h = _mm_xor_si128(prod_b_d, _mm_srli_si128::<8>(prod_ab_cd));

        let mut low = MaybeUninit::<[u64; 2]>::uninit();
        _mm_storeu_si128(low.as_mut_ptr() as *mut _, res_l);
        let mut high = MaybeUninit::<[u64; 2]>::uninit();
        _mm_storeu_si128(high.as_mut_ptr() as *mut _, res_h);

        (low.assume_init(), high.assume_init())
    }
}

#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
#[target_feature(enable = "pclmulqdq,sse2")]
unsafe fn clmul_u64_sse2(x: [u64; 1], y: [u64; 1]) -> ([u64; 1], [u64; 1]) {
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;
    use std::mem::MaybeUninit;

    let t = unsafe {
        let x = _mm_loadu_si64(x.as_ptr() as *const _);
        let y = _mm_loadu_si64(y.as_ptr() as *const _);

        let prod = _mm_clmulepi64_si128::<0>(x, y);

        let mut res = MaybeUninit::<[u64; 2]>::uninit();
        _mm_storeu_si128(res.as_mut_ptr() as *mut _, prod);

        res.assume_init()
    };

    ([t[0]], [t[1]])
}

// ARM NEON PMULL carry-less multiply (aarch64 with crypto extensions)

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon,aes")]
unsafe fn clmul_2_u64_neon(x: [u64; 2], y: [u64; 2]) -> ([u64; 2], [u64; 2]) {
    use core::arch::aarch64::{
        vgetq_lane_p64,
        vld1q_u64,
        vmull_high_p64,
        vmull_p64,
        vreinterpretq_p64_u64,
    };

    unsafe {
        let a = vreinterpretq_p64_u64(vld1q_u64(x.as_ptr()));
        let b = vreinterpretq_p64_u64(vld1q_u64(y.as_ptr()));

        // Karatsuba: lo = x[0]*y[0], hi = x[1]*y[1],
        // mid = (x[0]^x[1])*(y[0]^y[1])
        let lo = vmull_p64(vgetq_lane_p64::<0>(a), vgetq_lane_p64::<0>(b));
        let hi = vmull_high_p64(a, b);

        let s_x = x[0] ^ vgetq_lane_p64::<1>(a);
        let s_y = y[0] ^ vgetq_lane_p64::<1>(b);
        let mid = vmull_p64(s_x, s_y);

        let middle = mid ^ lo ^ hi;

        (
            [lo as u64, ((lo >> 64) ^ middle) as u64],
            [(hi ^ (middle >> 64)) as u64, (hi >> 64) as u64],
        )
    }
}

#[cfg(target_arch = "aarch64")]
#[target_feature(enable = "neon,aes")]
unsafe fn clmul_u64_neon(x: [u64; 1], y: [u64; 1]) -> ([u64; 1], [u64; 1]) {
    use core::arch::aarch64::vmull_p64;

    let prod = vmull_p64(x[0], y[0]);
    ([prod as u64], [(prod >> 64) as u64])
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_carry_less_mul_2limbs() {
        macro_rules! prod_test_case {
            ($aval:expr, $bval:expr, $prod:expr) => {{
                let prod_comp = carry_less_mul($aval, $bval);
                assert_eq!(prod_comp, $prod);
            }};
        }

        prod_test_case!(
            [0x9f418f3bffd84bba, 0x4a7c605645afdfb1],
            [0x80b7bd91cddc5be5, 0x3a97291035e41e1f],
            (
                [0xfffa315f244b1f92, 0x288de208c77fb2f9],
                [0x94aced70da538690, 0x0f0f341c05b65c5e,]
            )
        );

        prod_test_case!(
            [0x74ef862bc1b6d333, 0x3a88103b80d97b73],
            [0x753f4846eb020b5a, 0x8f108359ea25fa8f],
            (
                [0x419f60d2de880d0e, 0xfd4d74204161d27d],
                [0xe96db1bf781f351f, 0x1c36456adc21ac7a]
            )
        );

        prod_test_case!(
            [0x6447b3dcaed62649, 0x6e4af40b2ee1b4c1],
            [0xbd7a4e12fdb29840, 0x8950f56742015f25],
            (
                [0x010b1b56e559ca40, 0xf1ff4cfe33d20957],
                [0xed77c432e4701779, 0x342437199bebda57]
            )
        );
    }

    #[test]
    fn test_dispatched_matches_portable_2limbs() {
        let cases: [([u64; 2], [u64; 2]); 3] = [
            (
                [0x9f418f3bffd84bba, 0x4a7c605645afdfb1],
                [0x80b7bd91cddc5be5, 0x3a97291035e41e1f],
            ),
            (
                [0x74ef862bc1b6d333, 0x3a88103b80d97b73],
                [0x753f4846eb020b5a, 0x8f108359ea25fa8f],
            ),
            (
                [0x6447b3dcaed62649, 0x6e4af40b2ee1b4c1],
                [0xbd7a4e12fdb29840, 0x8950f56742015f25],
            ),
        ];
        for (a, b) in cases {
            let portable = carry_less_mul(a, b);
            let dispatched = carry_less_mul_2limbs(a, b);
            assert_eq!(portable, dispatched);
        }
    }

    #[test]
    fn test_dispatched_matches_portable_1limb() {
        let cases: [([u64; 1], [u64; 1]); 3] = [
            ([0x9f418f3bffd84bba], [0x80b7bd91cddc5be5]),
            ([0x74ef862bc1b6d333], [0x753f4846eb020b5a]),
            ([0xffffffffffffffff], [0xffffffffffffffff]),
        ];
        for (a, b) in cases {
            let portable = carry_less_mul(a, b);
            let dispatched = carry_less_mul_1limb(a, b);
            assert_eq!(portable, dispatched);
        }
    }
}