himada-dispatch 0.1.1

use himada_core::HardwareDNA;

// ---------------------------------------------------------------------------
// sort_i32
// ---------------------------------------------------------------------------

pub fn sort_i32_scalar(a: &mut [i32]) {
    a.sort_unstable();
}

pub fn sort_i32_supported(_: &HardwareDNA) -> bool { true }

#[cfg(target_arch = "x86_64")]
pub fn sort_i32_sse(a: &mut [i32]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;

    // SAFETY:
    // - `is_x86_feature_detected!("sse2")` confirms CPU support
    // - `i + 4 <= a.len()` ensures in-bounds access for `_mm_loadu_si128` and `_mm_storeu_si128`
    // - `a.as_ptr().add(i)` and `a.as_mut_ptr().add(i)` are valid because `i + 4 <= a.len()`
    // - `_mm_shuffle_epi32`, `_mm_min_epi32`, `_mm_max_epi32`, `_mm_unpacklo_epi32`,
    //   `_mm_unpackhi_epi32`, `_mm_castps_si128`, `_mm_castsi128_ps`, `_mm_movelh_ps`
    //   are pure compute intrinsics operating on already-loaded register values
    // - `a.sort_unstable()` after the SIMD loop ensures correctness for the remainder
    unsafe {
        if is_x86_feature_detected!("sse2") {
            let mut i = 0;
            while i + 4 <= a.len() {
                let v = _mm_loadu_si128(a.as_ptr().add(i) as *const __m128i);
                let s1 = _mm_shuffle_epi32(v, 0xB1);
                let lo1 = _mm_min_epi32(v, s1);
                let hi1 = _mm_max_epi32(v, s1);
                let u1 = _mm_unpacklo_epi32(lo1, hi1);
                let u2 = _mm_unpackhi_epi32(lo1, hi1);
                let v = _mm_castps_si128(_mm_movelh_ps(_mm_castsi128_ps(u1), _mm_castsi128_ps(u2)));
                let s2 = _mm_shuffle_epi32(v, 0x4E);
                let lo2 = _mm_min_epi32(v, s2);
                let hi2 = _mm_max_epi32(v, s2);
                let u1 = _mm_unpacklo_epi32(lo2, hi2);
                let u2 = _mm_unpackhi_epi32(lo2, hi2);
                let v = _mm_castps_si128(_mm_movelh_ps(_mm_castsi128_ps(u1), _mm_castsi128_ps(u2)));
                let s3 = _mm_shuffle_epi32(v, 0xB1);
                let lo3 = _mm_min_epi32(v, s3);
                let hi3 = _mm_max_epi32(v, s3);
                let u1 = _mm_unpacklo_epi32(lo3, hi3);
                let u2 = _mm_unpackhi_epi32(lo3, hi3);
                let v = _mm_castps_si128(_mm_movelh_ps(_mm_castsi128_ps(u1), _mm_castsi128_ps(u2)));
                _mm_storeu_si128(a.as_mut_ptr().add(i) as *mut __m128i, v);
                i += 4;
            }
        }
    }
    a.sort_unstable();
}

#[cfg(target_arch = "x86_64")]
pub fn sort_i32_sse_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "SSE2")
}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_i32_sse(_: &mut [i32]) {}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_i32_sse_supported(_: &HardwareDNA) -> bool { false }

#[cfg(target_arch = "x86_64")]
pub fn sort_i32_avx2(a: &mut [i32]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;

    // SAFETY:
    // - `is_x86_feature_detected!("avx2")` confirms CPU support
    // - `i + 8 <= a.len()` ensures in-bounds access for `_mm256_loadu_si256` and `_mm256_storeu_si256`
    // - `a.as_ptr().add(i)` and `a.as_mut_ptr().add(i)` are valid because `i + 8 <= a.len()`
    // - All shuffle/min/max/cast intrinsics operate on already-loaded register values
    // - `a.sort_unstable()` after the SIMD loop ensures correctness for the remainder
    unsafe {
        if is_x86_feature_detected!("avx2") {
            let mut i = 0;
            while i + 8 <= a.len() {
                let v = _mm256_loadu_si256(a.as_ptr().add(i) as *const __m256i);
                let s1 = _mm256_shuffle_epi32(v, 0xB1);
                let lo1 = _mm256_min_epi32(v, s1);
                let hi1 = _mm256_max_epi32(v, s1);
                let u1 = _mm256_unpacklo_epi32(lo1, hi1);
                let u2 = _mm256_unpackhi_epi32(lo1, hi1);
                let v = _mm256_castps_si256(_mm256_movelh_ps(_mm256_castsi256_ps(u1), _mm256_castsi256_ps(u2)));
                let s2 = _mm256_shuffle_epi32(v, 0x4E);
                let lo2 = _mm256_min_epi32(v, s2);
                let hi2 = _mm256_max_epi32(v, s2);
                let u1 = _mm256_unpacklo_epi32(lo2, hi2);
                let u2 = _mm256_unpackhi_epi32(lo2, hi2);
                let v = _mm256_castps_si256(_mm256_movelh_ps(_mm256_castsi256_ps(u1), _mm256_castsi256_ps(u2)));
                let s3 = _mm256_shuffle_epi32(v, 0xB1);
                let lo3 = _mm256_min_epi32(v, s3);
                let hi3 = _mm256_max_epi32(v, s3);
                let u1 = _mm256_unpacklo_epi32(lo3, hi3);
                let u2 = _mm256_unpackhi_epi32(lo3, hi3);
                let v = _mm256_castps_si256(_mm256_movelh_ps(_mm256_castsi256_ps(u1), _mm256_castsi256_ps(u2)));
                _mm256_storeu_si256(a.as_mut_ptr().add(i) as *mut __m256i, v);
                i += 8;
            }
        }
    }
    a.sort_unstable();
}

#[cfg(target_arch = "x86_64")]
pub fn sort_i32_avx2_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "AVX2")
}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_i32_avx2(_: &mut [i32]) {}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_i32_avx2_supported(_: &HardwareDNA) -> bool { false }

#[cfg(target_arch = "aarch64")]
pub fn sort_i32_neon(a: &mut [i32]) {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;

    // SAFETY:
    // - NEON is available on all aarch64 targets
    // - `i + 4 <= a.len()` ensures in-bounds access for `vld1q_s32` and `vst1q_s32`
    // - `a.as_ptr().add(i)` and `a.as_mut_ptr().add(i)` are valid because `i + 4 <= a.len()`
    // - All NEON intrinsics operate on already-loaded register values
    // - `a.sort_unstable()` after the SIMD loop ensures correctness for the remainder
    unsafe {
        let mut i = 0;
        while i + 4 <= a.len() {
            let v = vld1q_s32(a.as_ptr().add(i));
            let s1 = vrev64q_s32(v);
            let lo1 = vminq_s32(v, s1);
            let hi1 = vmaxq_s32(v, s1);
            let u1 = vtrn1q_s32(lo1, hi1);
            let u2 = vtrn2q_s32(lo1, hi1);
            let v = vuzp1q_s32(u1, u2);
            let s2 = vextq_s32(v, v, 2);
            let lo2 = vminq_s32(v, s2);
            let hi2 = vmaxq_s32(v, s2);
            let u1 = vtrn1q_s32(lo2, hi2);
            let u2 = vtrn2q_s32(lo2, hi2);
            let v = vuzp1q_s32(u1, u2);
            let s3 = vrev64q_s32(v);
            let lo3 = vminq_s32(v, s3);
            let hi3 = vmaxq_s32(v, s3);
            let u1 = vtrn1q_s32(lo3, hi3);
            let u2 = vtrn2q_s32(lo3, hi3);
            let v = vuzp1q_s32(u1, u2);
            vst1q_s32(a.as_mut_ptr().add(i), v);
            i += 4;
        }
    }
    a.sort_unstable();
}

#[cfg(target_arch = "aarch64")]
pub fn sort_i32_neon_supported(_: &HardwareDNA) -> bool { true }

#[cfg(not(target_arch = "aarch64"))]
pub fn sort_i32_neon(_: &mut [i32]) {}

#[cfg(not(target_arch = "aarch64"))]
pub fn sort_i32_neon_supported(_: &HardwareDNA) -> bool { false }

// ---------------------------------------------------------------------------
// sort_f64
// ---------------------------------------------------------------------------

pub fn sort_f64_scalar(a: &mut [f64]) {
    a.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
}

pub fn sort_f64_supported(_: &HardwareDNA) -> bool { true }

#[cfg(target_arch = "x86_64")]
pub fn sort_f64_sse(a: &mut [f64]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;

    // SAFETY:
    // - `is_x86_feature_detected!("sse2")` confirms CPU support
    // - `i + 2 <= a.len()` ensures in-bounds access for `_mm_loadu_pd` and `_mm_storeu_pd`
    // - `_mm_shuffle_pd`, `_mm_min_pd`, `_mm_max_pd`, `_mm_unpacklo_pd` are pure compute
    // - `a.sort_unstable_by(...)` after the SIMD loop ensures correctness for the remainder
    unsafe {
        if is_x86_feature_detected!("sse2") {
            let mut i = 0;
            while i + 2 <= a.len() {
                let v = _mm_loadu_pd(a.as_ptr().add(i));
                let swapped = _mm_shuffle_pd(v, v, 1);
                let lo = _mm_min_pd(v, swapped);
                let hi = _mm_max_pd(v, swapped);
                _mm_storeu_pd(a.as_mut_ptr().add(i), _mm_unpacklo_pd(lo, hi));
                i += 2;
            }
        }
    }
    a.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
}

#[cfg(target_arch = "x86_64")]
pub fn sort_f64_sse_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "SSE2")
}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_f64_sse(_: &mut [f64]) {}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_f64_sse_supported(_: &HardwareDNA) -> bool { false }

#[cfg(target_arch = "x86_64")]
pub fn sort_f64_avx2(a: &mut [f64]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;

    // SAFETY:
    // - `is_x86_feature_detected!("avx2")` confirms CPU support
    // - `i + 4 <= a.len()` ensures in-bounds access for `_mm256_loadu_pd` and `_mm256_storeu_pd`
    // - All shuffle/min/max/cast intrinsics operate on already-loaded register values
    // - `a.sort_unstable_by(...)` after the SIMD loop ensures correctness for the remainder
    unsafe {
        if is_x86_feature_detected!("avx2") {
            let mut i = 0;
            while i + 4 <= a.len() {
                let v = _mm256_loadu_pd(a.as_ptr().add(i));
                let s1 = _mm256_shuffle_pd(v, v, 5);
                let lo1 = _mm256_min_pd(v, s1);
                let hi1 = _mm256_max_pd(v, s1);
                let u1 = _mm256_unpacklo_pd(lo1, hi1);
                let u2 = _mm256_unpackhi_pd(lo1, hi1);
                let lo128 = _mm256_castpd256_pd128(u1);
                let hi128 = _mm256_castpd256_pd128(u2);
                let v = _mm256_insertf128_pd(_mm256_castpd128_pd256(lo128), hi128, 1);
                let s2 = _mm256_permute4x64_pd(v, 0x4E);
                let lo2 = _mm256_min_pd(v, s2);
                let hi2 = _mm256_max_pd(v, s2);
                let lo128_2 = _mm_unpacklo_pd(_mm256_castpd256_pd128(lo2), _mm256_castpd256_pd128(hi2));
                let hi128_2 = _mm_unpackhi_pd(_mm256_extractf128_pd(lo2, 1), _mm256_extractf128_pd(hi2, 1));
                let v = _mm256_insertf128_pd(_mm256_castpd128_pd256(lo128_2), hi128_2, 1);
                let s3 = _mm256_shuffle_pd(v, v, 5);
                let lo3 = _mm256_min_pd(v, s3);
                let hi3 = _mm256_max_pd(v, s3);
                let lo128_3 = _mm_unpacklo_pd(_mm256_castpd256_pd128(lo3), _mm256_castpd256_pd128(hi3));
                let hi128_3 = _mm_unpackhi_pd(_mm256_extractf128_pd(lo3, 1), _mm256_extractf128_pd(hi3, 1));
                let v = _mm256_insertf128_pd(_mm256_castpd128_pd256(lo128_3), hi128_3, 1);
                _mm256_storeu_pd(a.as_mut_ptr().add(i), v);
                i += 4;
            }
        }
    }
    a.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
}

#[cfg(target_arch = "x86_64")]
pub fn sort_f64_avx2_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "AVX2")
}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_f64_avx2(_: &mut [f64]) {}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_f64_avx2_supported(_: &HardwareDNA) -> bool { false }

#[cfg(target_arch = "aarch64")]
pub fn sort_f64_neon(a: &mut [f64]) {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;

    // SAFETY:
    // - NEON is available on all aarch64 targets
    // - `i + 2 <= a.len()` ensures in-bounds access for `vld1q_f64` and `vst1q_f64`
    // - All NEON intrinsics operate on already-loaded register values
    // - `a.sort_unstable_by(...)` after the SIMD loop ensures correctness for the remainder
    unsafe {
        let mut i = 0;
        while i + 2 <= a.len() {
            let v = vld1q_f64(a.as_ptr().add(i));
            let swapped = vextq_f64(v, v, 1);
            let lo = vminq_f64(v, swapped);
            let hi = vmaxq_f64(v, swapped);
            vst1q_f64(a.as_mut_ptr().add(i), vtrn1q_f64(lo, hi));
            i += 2;
        }
    }
    a.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
}

#[cfg(target_arch = "aarch64")]
pub fn sort_f64_neon_supported(_: &HardwareDNA) -> bool { true }

#[cfg(not(target_arch = "aarch64"))]
pub fn sort_f64_neon(_: &mut [f64]) {}

#[cfg(not(target_arch = "aarch64"))]
pub fn sort_f64_neon_supported(_: &HardwareDNA) -> bool { false }

// ---------------------------------------------------------------------------
// sort_f32
// ---------------------------------------------------------------------------

pub fn sort_f32_scalar(a: &mut [f32]) {
    a.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
}

pub fn sort_f32_supported(_: &HardwareDNA) -> bool { true }

#[cfg(target_arch = "x86_64")]
pub fn sort_f32_sse(a: &mut [f32]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;

    // SAFETY:
    // - `is_x86_feature_detected!("sse2")` confirms CPU support
    // - `i + 4 <= a.len()` ensures in-bounds access for `_mm_loadu_ps` and `_mm_storeu_ps`
    // - All shuffle/min/max intrinsics operate on already-loaded register values
    // - `a.sort_unstable_by(...)` after the SIMD loop ensures correctness for the remainder
    unsafe {
        if is_x86_feature_detected!("sse2") {
            let mut i = 0;
            while i + 4 <= a.len() {
                let v = _mm_loadu_ps(a.as_ptr().add(i));
                let s1 = _mm_shuffle_ps(v, v, 0xB1);
                let lo1 = _mm_min_ps(v, s1);
                let hi1 = _mm_max_ps(v, s1);
                let u1 = _mm_unpacklo_ps(lo1, hi1);
                let u2 = _mm_unpackhi_ps(lo1, hi1);
                let v = _mm_movelh_ps(u1, u2);
                let s2 = _mm_shuffle_ps(v, v, 0x4E);
                let lo2 = _mm_min_ps(v, s2);
                let hi2 = _mm_max_ps(v, s2);
                let u1 = _mm_unpacklo_ps(lo2, hi2);
                let u2 = _mm_unpackhi_ps(lo2, hi2);
                let v = _mm_movelh_ps(u1, u2);
                let s3 = _mm_shuffle_ps(v, v, 0xB1);
                let lo3 = _mm_min_ps(v, s3);
                let hi3 = _mm_max_ps(v, s3);
                let u1 = _mm_unpacklo_ps(lo3, hi3);
                let u2 = _mm_unpackhi_ps(lo3, hi3);
                let v = _mm_movelh_ps(u1, u2);
                _mm_storeu_ps(a.as_mut_ptr().add(i), v);
                i += 4;
            }
        }
    }
    a.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
}

#[cfg(target_arch = "x86_64")]
pub fn sort_f32_sse_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "SSE2")
}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_f32_sse(_: &mut [f32]) {}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_f32_sse_supported(_: &HardwareDNA) -> bool { false }

#[cfg(target_arch = "x86_64")]
pub fn sort_f32_avx2(a: &mut [f32]) {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;

    // SAFETY:
    // - `is_x86_feature_detected!("avx2")` confirms CPU support
    // - `i + 8 <= a.len()` ensures in-bounds access for `_mm256_loadu_ps` and `_mm256_storeu_ps`
    // - All shuffle/min/max intrinsics operate on already-loaded register values
    // - `a.sort_unstable_by(...)` after the SIMD loop ensures correctness for the remainder
    unsafe {
        if is_x86_feature_detected!("avx2") {
            let mut i = 0;
            while i + 8 <= a.len() {
                let v = _mm256_loadu_ps(a.as_ptr().add(i));
                let s1 = _mm256_shuffle_ps(v, v, 0xB1);
                let lo1 = _mm256_min_ps(v, s1);
                let hi1 = _mm256_max_ps(v, s1);
                let u1 = _mm256_unpacklo_ps(lo1, hi1);
                let u2 = _mm256_unpackhi_ps(lo1, hi1);
                let v = _mm256_movelh_ps(u1, u2);
                let s2 = _mm256_shuffle_ps(v, v, 0x4E);
                let lo2 = _mm256_min_ps(v, s2);
                let hi2 = _mm256_max_ps(v, s2);
                let u1 = _mm256_unpacklo_ps(lo2, hi2);
                let u2 = _mm256_unpackhi_ps(lo2, hi2);
                let v = _mm256_movelh_ps(u1, u2);
                let s3 = _mm256_shuffle_ps(v, v, 0xB1);
                let lo3 = _mm256_min_ps(v, s3);
                let hi3 = _mm256_max_ps(v, s3);
                let u1 = _mm256_unpacklo_ps(lo3, hi3);
                let u2 = _mm256_unpackhi_ps(lo3, hi3);
                let v = _mm256_movelh_ps(u1, u2);
                _mm256_storeu_ps(a.as_mut_ptr().add(i), v);
                i += 8;
            }
        }
    }
    a.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
}

#[cfg(target_arch = "x86_64")]
pub fn sort_f32_avx2_supported(dna: &HardwareDNA) -> bool {
    dna.cpu.features.iter().any(|f| f == "AVX2")
}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_f32_avx2(_: &mut [f32]) {}

#[cfg(not(target_arch = "x86_64"))]
pub fn sort_f32_avx2_supported(_: &HardwareDNA) -> bool { false }

#[cfg(target_arch = "aarch64")]
pub fn sort_f32_neon(a: &mut [f32]) {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;

    // SAFETY:
    // - NEON is available on all aarch64 targets
    // - `i + 4 <= a.len()` ensures in-bounds access for `vld1q_f32` and `vst1q_f32`
    // - All NEON intrinsics operate on already-loaded register values
    // - `a.sort_unstable_by(...)` after the SIMD loop ensures correctness for the remainder
    unsafe {
        let mut i = 0;
        while i + 4 <= a.len() {
            let v = vld1q_f32(a.as_ptr().add(i));
            let s1 = vrev64q_f32(v);
            let lo1 = vminq_f32(v, s1);
            let hi1 = vmaxq_f32(v, s1);
            let u1 = vtrn1q_f32(lo1, hi1);
            let u2 = vtrn2q_f32(lo1, hi1);
            let v = vuzp1q_f32(u1, u2);
            let s2 = vextq_f32(v, v, 2);
            let lo2 = vminq_f32(v, s2);
            let hi2 = vmaxq_f32(v, s2);
            let u1 = vtrn1q_f32(lo2, hi2);
            let u2 = vtrn2q_f32(lo2, hi2);
            let v = vuzp1q_f32(u1, u2);
            let s3 = vrev64q_f32(v);
            let lo3 = vminq_f32(v, s3);
            let hi3 = vmaxq_f32(v, s3);
            let u1 = vtrn1q_f32(lo3, hi3);
            let u2 = vtrn2q_f32(lo3, hi3);
            let v = vuzp1q_f32(u1, u2);
            vst1q_f32(a.as_mut_ptr().add(i), v);
            i += 4;
        }
    }
    a.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap());
}

#[cfg(target_arch = "aarch64")]
pub fn sort_f32_neon_supported(_: &HardwareDNA) -> bool { true }

#[cfg(not(target_arch = "aarch64"))]
pub fn sort_f32_neon(_: &mut [f32]) {}

#[cfg(not(target_arch = "aarch64"))]
pub fn sort_f32_neon_supported(_: &HardwareDNA) -> bool { false }