xocomil 0.3.0

//! SIMD-accelerated byte searching as an extension trait on `[u8]`.
//!
//! Uses SSE2 on x86-64, NEON on aarch64, and wasm SIMD on wasm32
//! (with `simd128`). Other architectures get a scalar fallback.

/// Extension trait providing fast byte-pattern searches on byte slices.
///
/// On x86-64 targets, searches use SSE2 intrinsics (guaranteed available
/// on all x86-64 processors). On aarch64, NEON intrinsics are used.
/// On wasm32 with `simd128`, wasm SIMD intrinsics are used.
/// Other architectures get a scalar fallback.
pub trait ByteSearch {
    /// Find the first occurrence of `needle`.
    fn find_byte(&self, needle: u8) -> Option<usize>;

    /// Find the first `\r\n` at or after `start`.
    fn find_crlf(&self, start: usize) -> Option<usize>;

    /// Find the `\r\n\r\n` header terminator at or after `start`.
    fn find_header_end(&self, start: usize) -> Option<usize>;
}

// ---------------------------------------------------------------------------
// Shared bitmask-based implementation (SSE2 + wasm SIMD)
// ---------------------------------------------------------------------------

/// Generates `find_byte`, `find_crlf`, and `find_header_end` using a
/// bitmask SIMD strategy. Each call site defines `simd_splat!`,
/// `simd_load!`, and `simd_mask!` helper macros that abstract over
/// the architecture-specific intrinsics before invoking this macro.
macro_rules! impl_bitmask_byte_search {
    () => {
        /// Safety: SIMD loads stay within `buf[i..i+16]` where `i + 16 <= len`.
        /// Scalar tail uses `get_unchecked(i)` where `i < len` (loop guard).
        #[inline]
        pub(super) fn find_byte(buf: &[u8], needle: u8) -> Option<usize> {
            let len = buf.len();
            unsafe {
                let nv = simd_splat!(needle);
                let mut i = 0;

                while i + 16 <= len {
                    let chunk = simd_load!(buf.as_ptr().add(i));
                    let mask = simd_mask!(chunk, nv);
                    if mask != 0 {
                        return Some(i + mask.trailing_zeros() as usize);
                    }
                    i += 16;
                }

                while i < len {
                    if *buf.get_unchecked(i) == needle {
                        return Some(i);
                    }
                    i += 1;
                }

                None
            }
        }

        /// Safety: SIMD loads stay within `buf[i..i+16]` where `i + 16 <= len`.
        /// `end = buf.len() - 1`, so `pos + 1` is always in bounds when `pos < end`.
        /// Scalar tail uses `get_unchecked(i)` and `get_unchecked(i+1)` where
        /// `i < end` (loop guard) and `end = len - 1`, so both accesses are in bounds.
        #[inline]
        pub(super) fn find_crlf(buf: &[u8], start: usize, end: usize) -> Option<usize> {
            let len = buf.len();
            unsafe {
                let cr = simd_splat!(HttpChar::CarriageReturn.as_u8());
                let mut i = start;

                while i + 16 <= len {
                    let chunk = simd_load!(buf.as_ptr().add(i));
                    let mut mask = simd_mask!(chunk, cr);

                    while mask != 0 {
                        let bit = mask.trailing_zeros() as usize;
                        let pos = i + bit;
                        if pos < end && *buf.get_unchecked(pos + 1) == HttpChar::LineFeed {
                            return Some(pos);
                        }
                        mask &= mask - 1;
                    }

                    i += 16;
                }

                while i < end {
                    if *buf.get_unchecked(i) == HttpChar::CarriageReturn
                        && *buf.get_unchecked(i + 1) == HttpChar::LineFeed
                    {
                        return Some(i);
                    }
                    i += 1;
                }

                None
            }
        }

        /// Safety: SIMD loads stay within `buf[i..i+16]` where `i + 16 <= len`.
        /// `end = buf.len() - 3`, so `pos + 1..pos + 3` is always in bounds
        /// when `pos < end`.
        /// Scalar tail uses `get_unchecked(i..i+3)` where `i < end` (loop guard)
        /// and `end = len - 3`, so all four accesses are in bounds.
        #[inline]
        pub(super) fn find_header_end(buf: &[u8], start: usize, end: usize) -> Option<usize> {
            let len = buf.len();
            unsafe {
                let cr = simd_splat!(HttpChar::CarriageReturn.as_u8());
                let mut i = start;

                while i + 16 <= len {
                    let chunk = simd_load!(buf.as_ptr().add(i));
                    let mut mask = simd_mask!(chunk, cr);

                    while mask != 0 {
                        let bit = mask.trailing_zeros() as usize;
                        let pos = i + bit;
                        if pos < end
                            && *buf.get_unchecked(pos + 1) == HttpChar::LineFeed
                            && *buf.get_unchecked(pos + 2) == HttpChar::CarriageReturn
                            && *buf.get_unchecked(pos + 3) == HttpChar::LineFeed
                        {
                            return Some(pos + 4);
                        }
                        mask &= mask - 1;
                    }

                    i += 16;
                }

                while i < end {
                    if *buf.get_unchecked(i) == HttpChar::CarriageReturn
                        && *buf.get_unchecked(i + 1) == HttpChar::LineFeed
                        && *buf.get_unchecked(i + 2) == HttpChar::CarriageReturn
                        && *buf.get_unchecked(i + 3) == HttpChar::LineFeed
                    {
                        return Some(i + 4);
                    }
                    i += 1;
                }

                None
            }
        }
    };
}

// ---------------------------------------------------------------------------
// Architecture-specific helpers
// ---------------------------------------------------------------------------

#[cfg(target_arch = "x86_64")]
#[allow(
    clippy::cast_possible_wrap,
    clippy::cast_sign_loss,
    clippy::cast_ptr_alignment
)]
mod sse2 {
    crate::simd::define_simd_primitives!();

    use crate::ascii::HttpChar;

    impl_bitmask_byte_search!();
}

#[cfg(target_arch = "aarch64")]
mod neon {
    use std::arch::aarch64::{
        vceqq_u8, vdupq_n_u8, vget_lane_u64, vld1q_u8, vreinterpret_u64_u8, vreinterpretq_u16_u8,
        vshrn_n_u16,
    };

    use crate::ascii::HttpChar;

    /// Extract a 16-bit bitmask from a NEON comparison result.
    ///
    /// Uses `vshrn_n_u16` to narrow each pair of bytes to a nibble,
    /// producing an 8-byte value where each matching lane contributes
    /// a `0x0F` nibble. The bit position `/ 4` gives the byte index.
    #[inline]
    unsafe fn neon_movemask(cmp: std::arch::aarch64::uint8x16_t) -> u64 {
        let narrowed = vshrn_n_u16(vreinterpretq_u16_u8(cmp), 4);
        vget_lane_u64(vreinterpret_u64_u8(narrowed), 0)
    }

    /// Safety: NEON is guaranteed on all aarch64 processors.
    /// SIMD loads stay within `buf[i..i+16]` where `i + 16 <= len`.
    /// Scalar tail uses `get_unchecked(i)` where `i < len` (loop guard).
    #[inline]
    pub(super) fn find_byte(buf: &[u8], needle: u8) -> Option<usize> {
        let len = buf.len();
        unsafe {
            let nv = vdupq_n_u8(needle);
            let mut i = 0;

            while i + 16 <= len {
                let chunk = vld1q_u8(buf.as_ptr().add(i));
                let cmp = vceqq_u8(chunk, nv);
                let bits = neon_movemask(cmp);
                if bits != 0 {
                    return Some(i + (bits.trailing_zeros() as usize) / 4);
                }
                i += 16;
            }

            while i < len {
                if *buf.get_unchecked(i) == needle {
                    return Some(i);
                }
                i += 1;
            }

            None
        }
    }

    /// Safety: NEON is guaranteed. SIMD loads stay within `buf[i..i+16]`
    /// where `i + 16 <= len`. `end = buf.len() - 1`, so `pos + 1` is in
    /// bounds when `pos < end`. Scalar tail: `get_unchecked(i)` and
    /// `get_unchecked(i+1)` where `i < end = len - 1`.
    #[inline]
    pub(super) fn find_crlf(buf: &[u8], start: usize, end: usize) -> Option<usize> {
        let len = buf.len();
        unsafe {
            let cr = vdupq_n_u8(HttpChar::CarriageReturn.as_u8());
            let mut i = start;

            while i + 16 <= len {
                let chunk = vld1q_u8(buf.as_ptr().add(i));
                let cmp = vceqq_u8(chunk, cr);
                let mut bits = neon_movemask(cmp);

                while bits != 0 {
                    let bit_pos = bits.trailing_zeros() as usize;
                    let pos = i + bit_pos / 4;
                    if pos < end && *buf.get_unchecked(pos + 1) == HttpChar::LineFeed {
                        return Some(pos);
                    }
                    // Clear this nibble (4 bits per lane).
                    // Each byte lane occupies 4 bits in the narrowed
                    // mask; clear the nibble for this lane and continue.
                    bits &= !(0xFu64 << (bit_pos & !3));
                }

                i += 16;
            }

            while i < end {
                if *buf.get_unchecked(i) == HttpChar::CarriageReturn
                    && *buf.get_unchecked(i + 1) == HttpChar::LineFeed
                {
                    return Some(i);
                }
                i += 1;
            }

            None
        }
    }

    /// Safety: NEON is guaranteed. SIMD loads stay within `buf[i..i+16]`
    /// where `i + 16 <= len`. `end = buf.len() - 3`, so `pos + 1..pos + 3`
    /// is in bounds when `pos < end`. Scalar tail: `get_unchecked(i..i+3)`
    /// where `i < end = len - 3`.
    #[inline]
    pub(super) fn find_header_end(buf: &[u8], start: usize, end: usize) -> Option<usize> {
        let len = buf.len();
        unsafe {
            let cr = vdupq_n_u8(HttpChar::CarriageReturn.as_u8());
            let mut i = start;

            while i + 16 <= len {
                let chunk = vld1q_u8(buf.as_ptr().add(i));
                let cmp = vceqq_u8(chunk, cr);
                let mut bits = neon_movemask(cmp);

                while bits != 0 {
                    let bit_pos = bits.trailing_zeros() as usize;
                    let pos = i + bit_pos / 4;
                    if pos < end
                        && *buf.get_unchecked(pos + 1) == HttpChar::LineFeed
                        && *buf.get_unchecked(pos + 2) == HttpChar::CarriageReturn
                        && *buf.get_unchecked(pos + 3) == HttpChar::LineFeed
                    {
                        return Some(pos + 4);
                    }
                    // Each byte lane occupies 4 bits in the narrowed
                    // mask; clear the nibble for this lane and continue.
                    bits &= !(0xFu64 << (bit_pos & !3));
                }

                i += 16;
            }

            while i < end {
                if *buf.get_unchecked(i) == HttpChar::CarriageReturn
                    && *buf.get_unchecked(i + 1) == HttpChar::LineFeed
                    && *buf.get_unchecked(i + 2) == HttpChar::CarriageReturn
                    && *buf.get_unchecked(i + 3) == HttpChar::LineFeed
                {
                    return Some(i + 4);
                }
                i += 1;
            }

            None
        }
    }
}

#[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
#[allow(clippy::cast_sign_loss)]
mod wasm_simd {
    crate::simd::define_simd_primitives!();

    use crate::ascii::HttpChar;

    impl_bitmask_byte_search!();
}

#[cfg(not(any(
    target_arch = "x86_64",
    target_arch = "aarch64",
    all(target_arch = "wasm32", target_feature = "simd128")
)))]
mod scalar {
    use crate::ascii::HttpChar;

    #[inline]
    pub(super) fn find_byte(buf: &[u8], needle: u8) -> Option<usize> {
        buf.iter().position(|&b| b == needle)
    }

    #[inline]
    pub(super) fn find_crlf(buf: &[u8], start: usize, end: usize) -> Option<usize> {
        let mut i = start;
        while i < end {
            if buf[i] == HttpChar::CarriageReturn && buf[i + 1] == HttpChar::LineFeed {
                return Some(i);
            }
            i += 1;
        }
        None
    }

    #[inline]
    pub(super) fn find_header_end(buf: &[u8], start: usize, end: usize) -> Option<usize> {
        let mut i = start;
        while i < end {
            if buf[i] == HttpChar::CarriageReturn
                && buf[i + 1] == HttpChar::LineFeed
                && buf[i + 2] == HttpChar::CarriageReturn
                && buf[i + 3] == HttpChar::LineFeed
            {
                return Some(i + 4);
            }
            i += 1;
        }
        None
    }
}

// ---------------------------------------------------------------------------
// Dispatch
// ---------------------------------------------------------------------------

/// Select the best implementation for the current target.
macro_rules! dispatch {
    ($fn_name:ident $(, $arg:expr)* $(,)?) => {{
        #[cfg(target_arch = "x86_64")]
        { sse2::$fn_name($($arg),*) }
        #[cfg(target_arch = "aarch64")]
        { neon::$fn_name($($arg),*) }
        #[cfg(all(target_arch = "wasm32", target_feature = "simd128"))]
        { wasm_simd::$fn_name($($arg),*) }
        #[cfg(not(any(
            target_arch = "x86_64",
            target_arch = "aarch64",
            all(target_arch = "wasm32", target_feature = "simd128")
        )))]
        { scalar::$fn_name($($arg),*) }
    }};
}

impl ByteSearch for [u8] {
    #[inline]
    fn find_byte(&self, needle: u8) -> Option<usize> {
        dispatch!(find_byte, self, needle)
    }

    #[inline]
    fn find_crlf(&self, start: usize) -> Option<usize> {
        let len = self.len();
        if len < 2 {
            return None;
        }
        dispatch!(find_crlf, self, start, len - 1)
    }

    #[inline]
    fn find_header_end(&self, start: usize) -> Option<usize> {
        let len = self.len();
        if len < 4 {
            return None;
        }
        dispatch!(find_header_end, self, start, len - 3)
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

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

    // -- find_byte -----------------------------------------------------------

    #[test]
    fn find_byte_empty() {
        assert_eq!(b"".find_byte(b'x'), None);
    }

    #[test]
    fn find_byte_single_match() {
        assert_eq!(b"x".find_byte(b'x'), Some(0));
    }

    #[test]
    fn find_byte_single_no_match() {
        assert_eq!(b"a".find_byte(b'x'), None);
    }

    #[test]
    fn find_byte_at_simd_boundary() {
        // Exactly 16 bytes, needle at each position
        for pos in 0..16 {
            let mut buf = [b'.'; 16];
            buf[pos] = b'x';
            assert_eq!(buf.find_byte(b'x'), Some(pos), "pos={pos}");
        }
    }

    #[test]
    fn find_byte_in_scalar_tail() {
        // 17 bytes: first 16 are misses, needle at position 16
        let mut buf = [b'.'; 17];
        buf[16] = b'x';
        assert_eq!(buf.find_byte(b'x'), Some(16));
    }

    #[test]
    fn find_byte_across_two_simd_chunks() {
        // 33 bytes, needle in second chunk
        let mut buf = [b'.'; 33];
        buf[20] = b'x';
        assert_eq!(buf.find_byte(b'x'), Some(20));
    }

    #[test]
    fn find_byte_no_match_large() {
        let buf = [b'.'; 64];
        assert_eq!(buf.find_byte(b'x'), None);
    }

    // -- find_crlf -----------------------------------------------------------

    #[test]
    fn find_crlf_empty() {
        assert_eq!(b"".find_crlf(0), None);
    }

    #[test]
    fn find_crlf_single_byte() {
        assert_eq!(b"\r".find_crlf(0), None);
    }

    #[test]
    fn find_crlf_at_start() {
        assert_eq!(b"\r\nrest".find_crlf(0), Some(0));
    }

    #[test]
    fn find_crlf_at_simd_boundary() {
        // CRLF at positions 14-15 (end of first 16-byte chunk)
        let mut buf = [b'.'; 16];
        buf[14] = b'\r';
        buf[15] = b'\n';
        assert_eq!(buf.find_crlf(0), Some(14));
    }

    #[test]
    fn find_crlf_spanning_simd_boundary() {
        // CR at position 15, LF at position 16
        let mut buf = [b'.'; 17];
        buf[15] = b'\r';
        buf[16] = b'\n';
        assert_eq!(buf.find_crlf(0), Some(15));
    }

    #[test]
    fn find_crlf_in_scalar_tail() {
        let mut buf = [b'.'; 19];
        buf[17] = b'\r';
        buf[18] = b'\n';
        assert_eq!(buf.find_crlf(0), Some(17));
    }

    #[test]
    fn find_crlf_bare_cr_no_lf() {
        let mut buf = [b'.'; 5];
        buf[2] = b'\r';
        assert_eq!(buf.find_crlf(0), None);
    }

    #[test]
    fn find_crlf_with_start_offset() {
        let buf = b"first\r\nsecond\r\n";
        assert_eq!(buf.find_crlf(0), Some(5));
        assert_eq!(buf.find_crlf(6), Some(13));
    }

    #[test]
    fn find_crlf_dense_cr_no_lf() {
        // Two full SIMD chunks of CR with no LF anywhere.
        // Exercises the per-lane nibble-clearing loop in the NEON path:
        // every byte matches the needle so the bitmask is fully set
        // (16 lanes / chunk), and the inner `while bits != 0` loop must
        // clear all 16 nibbles without infinite-looping.
        let buf = [b'\r'; 32];
        assert_eq!(buf.find_crlf(0), None);
    }

    #[test]
    fn find_crlf_dense_cr_with_one_lf() {
        // 33 bytes: positions 0..=31 are CR, position 32 is LF.
        // Both 16-byte SIMD chunks produce a fully-set bitmask, but the
        // only real CRLF is at position 31 (CR at 31, LF at 32). The
        // nibble-clearing loop must reject 15 false positives in the
        // second chunk before accepting the last lane.
        let mut buf = [b'\r'; 33];
        buf[32] = b'\n';
        assert_eq!(buf.find_crlf(0), Some(31));
    }

    // -- find_header_end -----------------------------------------------------

    #[test]
    fn find_header_end_empty() {
        assert_eq!(b"".find_header_end(0), None);
    }

    #[test]
    fn find_header_end_too_short() {
        assert_eq!(b"\r\n\r".find_header_end(0), None);
    }

    #[test]
    fn find_header_end_exact() {
        assert_eq!(b"\r\n\r\n".find_header_end(0), Some(4));
    }

    #[test]
    fn find_header_end_after_headers() {
        let buf = b"Host: localhost\r\n\r\nbody";
        assert_eq!(buf.find_header_end(0), Some(19));
    }

    #[test]
    fn find_header_end_at_simd_boundary() {
        // \r\n\r\n starting at position 12 (within first 16 bytes)
        let mut buf = [b'.'; 20];
        buf[12] = b'\r';
        buf[13] = b'\n';
        buf[14] = b'\r';
        buf[15] = b'\n';
        assert_eq!(buf.find_header_end(0), Some(16));
    }

    #[test]
    fn find_header_end_spanning_simd_boundary() {
        // \r\n\r\n starting at position 14 (spans into scalar tail)
        let mut buf = [b'.'; 20];
        buf[14] = b'\r';
        buf[15] = b'\n';
        buf[16] = b'\r';
        buf[17] = b'\n';
        assert_eq!(buf.find_header_end(0), Some(18));
    }

    #[test]
    fn find_header_end_in_scalar_tail() {
        let mut buf = [b'.'; 21];
        buf[17] = b'\r';
        buf[18] = b'\n';
        buf[19] = b'\r';
        buf[20] = b'\n';
        assert_eq!(buf.find_header_end(0), Some(21));
    }

    #[test]
    fn find_header_end_with_start_offset() {
        let buf = b"GET / HTTP/1.1\r\nHost: h\r\n\r\n";
        assert_eq!(buf.find_header_end(0), Some(27));
        // Starting after "Host" header line still finds it
        assert_eq!(buf.find_header_end(16), Some(27));
    }

    #[test]
    fn find_header_end_no_match() {
        let buf = b"GET / HTTP/1.1\r\nHost: h\r\n";
        assert_eq!(buf.find_header_end(0), None);
    }

    #[test]
    fn find_header_end_many_crs_no_terminator() {
        // Four full SIMD chunks of CR with no LF anywhere.
        // Every lane in every chunk matches the needle, producing a
        // fully-set 64-bit bitmask per chunk. The nibble-clearing loop
        // must reject all 16 false positives per chunk and terminate
        // without finding `\r\n\r\n`.
        let buf = [b'\r'; 64];
        assert_eq!(buf.find_header_end(0), None);
    }

    #[test]
    fn find_header_end_dense_crs_with_terminator() {
        // 64 bytes that are mostly CRs, with a real `\r\n\r\n` at
        // positions 30..=33. Because there are no LFs other than at
        // 31 and 33, no CR other than the one at 30 can pass the
        // four-byte LF-CR-LF tail check. This stresses the nibble
        // clearing path: each SIMD chunk has many lanes set, but only
        // one lane in the second chunk actually terminates the search.
        let mut buf = [b'\r'; 64];
        buf[31] = b'\n';
        buf[33] = b'\n';
        // Sanity: positions 30..=33 spell "\r\n\r\n"
        assert_eq!(&buf[30..34], b"\r\n\r\n");
        assert_eq!(buf.find_header_end(0), Some(34));
    }
}