xbbg_core 1.1.2

//! SIMD-accelerated utilities for xbbg-core
//!
//! Provides vectorized operations for:
//! - Validity bitmap packing (8-10x faster)
//! - ASCII detection for fast-path string handling (~20ns saved per string)
//! - Bulk numeric conversions (i32 → f64)
//!
//! All functions have scalar fallbacks when SIMD is unavailable.

// ============================================================================
// VALIDITY BITMAP PACKING
// ============================================================================

/// Pack validity bytes [u8; N] (0 or non-zero) into a bitmap (1 bit per row).
///
/// # Safety
/// - `out` must have capacity for at least `(valid.len() + 7) / 8` bytes
#[inline]
pub fn pack_validity(valid: &[u8], out: &mut [u8]) {
    #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
    {
        // SAFETY: We've verified AVX2 is available at compile time
        unsafe { pack_validity_avx2(valid, out) }
    }

    #[cfg(not(all(target_arch = "x86_64", target_feature = "avx2")))]
    {
        pack_validity_scalar(valid, out)
    }
}

/// Scalar implementation - processes 1 bit at a time
#[inline]
pub fn pack_validity_scalar(valid: &[u8], out: &mut [u8]) {
    // Zero the output first
    out.iter_mut().for_each(|b| *b = 0);

    for (i, &v) in valid.iter().enumerate() {
        if v != 0 {
            out[i / 8] |= 1 << (i % 8);
        }
    }
}

/// AVX2 implementation - processes 32 rows in ~3 cycles
///
/// # Safety
///
/// Caller must ensure the CPU supports AVX2 (check with `is_x86_feature_detected!("avx2")`).
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
pub unsafe fn pack_validity_avx2(valid: &[u8], out: &mut [u8]) {
    use std::arch::x86_64::*;

    let ptr = valid.as_ptr();
    let out_ptr = out.as_mut_ptr();
    let n = valid.len();

    let mut i = 0;
    let mut out_idx = 0;

    // Process 32 bytes at a time
    while i + 32 <= n {
        // Load 32 validity bytes
        let v = _mm256_loadu_si256(ptr.add(i) as *const __m256i);

        // Compare each byte to zero: 0xFF where zero, 0x00 where non-zero
        let is_zero = _mm256_cmpeq_epi8(v, _mm256_setzero_si256());

        // Extract MSB of each byte → 32 bits
        let mask = _mm256_movemask_epi8(is_zero) as u32;

        // Invert: we want 1 for valid (non-zero), not for zero
        let valid_mask = !mask;

        // Store 4 bytes (32 bits)
        (out_ptr.add(out_idx) as *mut u32).write_unaligned(valid_mask);

        out_idx += 4;
        i += 32;
    }

    // Scalar tail for remaining elements
    for j in i..n {
        if *ptr.add(j) != 0 {
            *out_ptr.add(j / 8) |= 1 << (j % 8);
        }
    }
}

// ============================================================================
// ASCII DETECTION
// ============================================================================

/// Check if a byte slice is pure ASCII (all bytes < 128).
///
/// Uses SIMD when available for 10x+ speedup on longer strings.
#[inline]
pub fn is_ascii(data: &[u8]) -> bool {
    #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
    {
        // SAFETY: We've verified AVX2 is available at compile time
        unsafe { is_ascii_avx2(data) }
    }

    #[cfg(not(all(target_arch = "x86_64", target_feature = "avx2")))]
    {
        is_ascii_scalar(data)
    }
}

/// Scalar ASCII check
#[inline]
pub fn is_ascii_scalar(data: &[u8]) -> bool {
    data.iter().all(|&b| b < 128)
}

/// AVX2 ASCII check - processes 32 bytes in ~3 cycles
///
/// # Safety
///
/// Caller must ensure the CPU supports AVX2 (check with `is_x86_feature_detected!("avx2")`).
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
pub unsafe fn is_ascii_avx2(data: &[u8]) -> bool {
    use std::arch::x86_64::*;

    let ptr = data.as_ptr();
    let len = data.len();
    let mut i = 0;

    // Check 32 bytes at a time
    while i + 32 <= len {
        let chunk = _mm256_loadu_si256(ptr.add(i) as *const __m256i);

        // movemask extracts MSB of each byte
        // ASCII has MSB = 0, so mask should be 0 for all-ASCII
        if _mm256_movemask_epi8(chunk) != 0 {
            return false; // Found non-ASCII byte
        }
        i += 32;
    }

    // Scalar check for tail
    data[i..].iter().all(|&b| b < 128)
}

/// Runtime AVX2 detection + dispatch for ASCII check
#[cfg(target_arch = "x86_64")]
#[inline]
pub fn is_ascii_runtime(data: &[u8]) -> bool {
    if is_x86_feature_detected!("avx2") {
        // SAFETY: We just verified AVX2 is available
        unsafe { is_ascii_avx2(data) }
    } else {
        is_ascii_scalar(data)
    }
}

#[cfg(not(target_arch = "x86_64"))]
#[inline]
pub fn is_ascii_runtime(data: &[u8]) -> bool {
    is_ascii_scalar(data)
}

// ============================================================================
// BULK NUMERIC CONVERSIONS
// ============================================================================

/// Convert i32 slice to f64 slice using SIMD when available.
///
/// # Panics
/// Panics if `dst.len() < src.len()`
#[inline]
pub fn i32_to_f64(src: &[i32], dst: &mut [f64]) {
    assert!(dst.len() >= src.len(), "destination too small");

    #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
    {
        // SAFETY: We've verified AVX2 is available at compile time
        unsafe { i32_to_f64_avx2(src, dst) }
    }

    #[cfg(not(all(target_arch = "x86_64", target_feature = "avx2")))]
    {
        i32_to_f64_scalar(src, dst)
    }
}

/// Scalar i32 to f64 conversion
#[inline]
pub fn i32_to_f64_scalar(src: &[i32], dst: &mut [f64]) {
    for (d, &s) in dst.iter_mut().zip(src.iter()) {
        *d = s as f64;
    }
}

/// AVX2 i32 to f64 conversion - processes 4 values at a time
///
/// # Safety
///
/// Caller must ensure the CPU supports AVX2 (check with `is_x86_feature_detected!("avx2")`).
#[cfg(target_arch = "x86_64")]
#[target_feature(enable = "avx2")]
#[inline]
pub unsafe fn i32_to_f64_avx2(src: &[i32], dst: &mut [f64]) {
    use std::arch::x86_64::*;

    let src_ptr = src.as_ptr();
    let dst_ptr = dst.as_mut_ptr();
    let len = src.len();
    let mut i = 0;

    // Process 4 i32s at a time → 4 f64s
    while i + 4 <= len {
        // Load 4 i32s (128 bits)
        let ints = _mm_loadu_si128(src_ptr.add(i) as *const __m128i);

        // Convert to 4 f64s (256 bits)
        let floats = _mm256_cvtepi32_pd(ints);

        // Store 4 f64s
        _mm256_storeu_pd(dst_ptr.add(i), floats);

        i += 4;
    }

    // Scalar tail
    for j in i..len {
        *dst_ptr.add(j) = *src_ptr.add(j) as f64;
    }
}

/// Runtime AVX2 detection + dispatch for i32→f64 conversion
#[cfg(target_arch = "x86_64")]
#[inline]
pub fn i32_to_f64_runtime(src: &[i32], dst: &mut [f64]) {
    assert!(dst.len() >= src.len(), "destination too small");

    if is_x86_feature_detected!("avx2") {
        // SAFETY: We just verified AVX2 is available
        unsafe { i32_to_f64_avx2(src, dst) }
    } else {
        i32_to_f64_scalar(src, dst)
    }
}

#[cfg(not(target_arch = "x86_64"))]
#[inline]
pub fn i32_to_f64_runtime(src: &[i32], dst: &mut [f64]) {
    i32_to_f64_scalar(src, dst)
}

// ============================================================================
// TESTS
// ============================================================================

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

    #[test]
    fn test_pack_validity_scalar() {
        // 8 values: [1, 0, 1, 1, 0, 0, 1, 0] → 0b01001101 = 0x4D
        let valid = [1u8, 0, 1, 1, 0, 0, 1, 0];
        let mut out = [0u8; 1];
        pack_validity_scalar(&valid, &mut out);
        assert_eq!(out[0], 0b01001101);
    }

    #[test]
    fn test_pack_validity_16() {
        // 16 values across 2 bytes
        let valid = [1u8, 1, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0, 1, 1, 1, 1];
        let mut out = [0u8; 2];
        pack_validity_scalar(&valid, &mut out);
        assert_eq!(out[0], 0b11110011); // bits 0-7
        assert_eq!(out[1], 0b11110000); // bits 8-15
    }

    #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
    #[test]
    fn test_pack_validity_avx2() {
        // 32 values
        let mut valid = [0u8; 32];
        valid[0] = 1;
        valid[7] = 1;
        valid[8] = 1;
        valid[31] = 1;

        let mut out_scalar = [0u8; 4];
        let mut out_avx2 = [0u8; 4];

        pack_validity_scalar(&valid, &mut out_scalar);
        unsafe { pack_validity_avx2(&valid, &mut out_avx2) };

        assert_eq!(out_scalar, out_avx2);
    }

    #[test]
    fn test_is_ascii_scalar() {
        assert!(is_ascii_scalar(b"hello world"));
        assert!(is_ascii_scalar(b"IBM US Equity"));
        assert!(is_ascii_scalar(b"PX_LAST"));
        assert!(!is_ascii_scalar("café".as_bytes()));
        assert!(!is_ascii_scalar(&[0x80]));
    }

    #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
    #[test]
    fn test_is_ascii_avx2() {
        // Test with >32 byte strings
        let ascii_long = b"This is a long ASCII string that exceeds 32 bytes easily!";
        let non_ascii = "This has a café in it somewhere in the middle here".as_bytes();

        unsafe {
            assert!(is_ascii_avx2(ascii_long));
            assert!(!is_ascii_avx2(non_ascii));
        }
    }

    #[test]
    fn test_i32_to_f64_scalar() {
        let src = [1i32, -2, 3, 1000000];
        let mut dst = [0.0f64; 4];
        i32_to_f64_scalar(&src, &mut dst);
        assert_eq!(dst, [1.0, -2.0, 3.0, 1000000.0]);
    }

    #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
    #[test]
    fn test_i32_to_f64_avx2() {
        let src: Vec<i32> = (0..100).collect();
        let mut dst_scalar = vec![0.0f64; 100];
        let mut dst_avx2 = vec![0.0f64; 100];

        i32_to_f64_scalar(&src, &mut dst_scalar);
        unsafe { i32_to_f64_avx2(&src, &mut dst_avx2) };

        assert_eq!(dst_scalar, dst_avx2);
    }
}