himada-dispatch 0.1.0

use himada_core::HardwareDNA;

/// Find first occurrence of `byte` in `data`. Returns `None` if not found.
pub fn memchr_scalar(byte: u8, data: &[u8]) -> Option<usize> {
    data.iter().position(|&b| b == byte)
}

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

#[cfg(target_arch = "x86_64")]
pub fn memchr_sse(byte: u8, data: &[u8]) -> Option<usize> {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let ptr = data.as_ptr();
    let len = data.len();
    let mut i = 0;

    // SAFETY:
    // - `is_x86_feature_detected!("sse2")` confirms CPU support for SSE2 intrinsics
    // - `len >= 16` ensures at least one full 16-byte chunk can be processed
    // - `ptr.add(i)` is in-bounds because `i + 16 <= len`
    // - `_mm_loadu_si128` works on unaligned pointers
    // - `_mm_set1_epi8`, `_mm_cmpeq_epi8`, `_mm_movemask_epi8` are pure compute intrinsics
    // - `mask.trailing_zeros()` is well-defined when `mask != 0` (guaranteed by the check)
    // - remainder is handled by scalar code after the block
    unsafe {
        if is_x86_feature_detected!("sse2") && len >= 16 {
            let vpat = _mm_set1_epi8(byte as i8);
            while i + 16 <= len {
                let chunk = _mm_loadu_si128(ptr.add(i) as *const __m128i);
                let cmp = _mm_cmpeq_epi8(chunk, vpat);
                let mask = _mm_movemask_epi8(cmp);
                if mask != 0 {
                    return Some(i + (mask.trailing_zeros() as usize));
                }
                i += 16;
            }
        }
    }

    if let Some(offset) = data[i..].iter().position(|&b| b == byte) {
        return Some(i + offset);
    }
    None
}

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

#[cfg(not(target_arch = "x86_64"))]
pub fn memchr_sse(_: u8, _: &[u8]) -> Option<usize> { None }

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

#[cfg(target_arch = "x86_64")]
pub fn memchr_avx2(byte: u8, data: &[u8]) -> Option<usize> {
    #[cfg(target_arch = "x86_64")]
    use std::arch::x86_64::*;
    let ptr = data.as_ptr();
    let len = data.len();
    let mut i = 0;

    // SAFETY:
    // - `is_x86_feature_detected!("avx2")` confirms CPU support for AVX2 intrinsics
    // - `len >= 32` ensures at least one full 32-byte chunk can be processed
    // - `ptr.add(i)` is in-bounds because `i + 32 <= len`
    // - `_mm256_loadu_si256` works on unaligned pointers
    // - `_mm256_set1_epi8`, `_mm256_cmpeq_epi8`, `_mm256_movemask_epi8` are pure compute intrinsics
    // - `mask.trailing_zeros()` is well-defined when `mask != 0` (guaranteed by the check)
    // - remainder is handled by scalar code after the block
    unsafe {
        if is_x86_feature_detected!("avx2") && len >= 32 {
            let vpat = _mm256_set1_epi8(byte as i8);
            while i + 32 <= len {
                let chunk = _mm256_loadu_si256(ptr.add(i) as *const __m256i);
                let cmp = _mm256_cmpeq_epi8(chunk, vpat);
                let mask = _mm256_movemask_epi8(cmp);
                if mask != 0 {
                    return Some(i + (mask.trailing_zeros() as usize));
                }
                i += 32;
            }
        }
    }

    if let Some(offset) = data[i..].iter().position(|&b| b == byte) {
        return Some(i + offset);
    }
    None
}

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

#[cfg(not(target_arch = "x86_64"))]
pub fn memchr_avx2(_: u8, _: &[u8]) -> Option<usize> { None }

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

#[cfg(target_arch = "aarch64")]
pub fn memchr_neon(byte: u8, data: &[u8]) -> Option<usize> {
    #[cfg(target_arch = "aarch64")]
    use std::arch::aarch64::*;
    let ptr = data.as_ptr();
    let len = data.len();
    let mut i = 0;

    // SAFETY:
    // - NEON is always available on aarch64
    // - `len >= 16` ensures at least one full 16-byte chunk can be processed
    // - `ptr.add(i)` is in-bounds because `i + 16 <= len`
    // - `vld1q_u8` tolerates unaligned pointers on aarch64
    // - `vdupq_n_u8`, `vceqq_u8`, `vreinterpretq_u64_u8`, `vgetq_lane_u64` are pure compute intrinsics
    // - `trailing_zeros()` on a non-zero mask always returns a valid index
    // - remainder is handled by scalar code after the block
    unsafe {
        if len >= 16 {
            let vpat = vdupq_n_u8(byte);
            while i + 16 <= len {
                let chunk = vld1q_u8(ptr.add(i));
                let cmp = vceqq_u8(chunk, vpat);
                let cmp64 = vreinterpretq_u64_u8(cmp);
                let mask_lo = vgetq_lane_u64(cmp64, 0);
                let mask_hi = vgetq_lane_u64(cmp64, 1);
                if mask_lo != 0 {
                    return Some(i + (mask_lo.trailing_zeros() as usize) / 8);
                }
                if mask_hi != 0 {
                    return Some(i + 8 + (mask_hi.trailing_zeros() as usize) / 8);
                }
                i += 16;
            }
        }
    }

    if let Some(offset) = data[i..].iter().position(|&b| b == byte) {
        return Some(i + offset);
    }
    None
}

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

#[cfg(not(target_arch = "aarch64"))]
pub fn memchr_neon(_: u8, _: &[u8]) -> Option<usize> { None }

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