ntex-httparse 2.2.0

use crate::iter::Bytes;

#[inline]
#[target_feature(enable = "avx2")]
pub unsafe fn match_header_name_vectored(bytes: &mut Bytes) {
    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;

    // see `sse42::match_header_name_char_16_sse`, `vpshufb` works on each
    // 128-bit lane separately, so the tables are duplicated in both lanes
    let bitmap = _mm_loadu_si128(crate::utils::TOKEN_NIBBLES.as_ptr() as *const _);
    let bits = _mm_setr_epi8(1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128);

    while bytes.as_ref().len() >= 32 {
        let bitmap = _mm256_broadcastsi128_si256(bitmap);
        let bits = _mm256_broadcastsi128_si256(bits);
        let dat = _mm256_lddqu_si256(bytes.as_ref().as_ptr() as *const _);
        let rows = _mm256_shuffle_epi8(
            bitmap,
            _mm256_and_si256(dat, _mm256_set1_epi8(0x8f_u8.cast_signed())),
        );
        let hi = _mm256_and_si256(_mm256_srli_epi16(dat, 4), _mm256_set1_epi8(0x0f));
        let bit = _mm256_shuffle_epi8(bits, hi);
        let ok = _mm256_cmpeq_epi8(_mm256_and_si256(rows, bit), bit);
        let advance = (_mm256_movemask_epi8(ok) as u32).trailing_ones() as usize;
        bytes.advance(advance);

        if advance != 32 {
            return;
        }
    }

    // header names are short, check 16 bytes before falling back to SWAR
    if bytes.as_ref().len() >= 16 {
        let dat = _mm_lddqu_si128(bytes.as_ref().as_ptr() as *const _);
        let rows = _mm_shuffle_epi8(
            bitmap,
            _mm_and_si128(dat, _mm_set1_epi8(0x8f_u8.cast_signed())),
        );
        let hi = _mm_and_si128(_mm_srli_epi16(dat, 4), _mm_set1_epi8(0x0f));
        let bit = _mm_shuffle_epi8(bits, hi);
        let ok = _mm_cmpeq_epi8(_mm_and_si128(rows, bit), bit);
        let advance = (_mm_movemask_epi8(ok) as u16).trailing_ones() as usize;
        bytes.advance(advance);

        if advance != 16 {
            return;
        }
    }
    super::swar::match_header_name_vectored(bytes);
}

#[inline]
#[target_feature(enable = "avx2")]
pub unsafe fn match_uri_vectored(bytes: &mut Bytes) {
    while bytes.as_ref().len() >= 32 {
        let advance = match_url_char_32_avx(bytes.as_ref());

        bytes.advance(advance);

        if advance != 32 {
            return;
        }
    }
    // NOTE: use SWAR for <32B, more efficient than falling back to SSE4.2
    super::swar::match_uri_vectored(bytes);
}

#[inline(always)]
#[allow(non_snake_case, overflowing_literals)]
#[allow(unused)]
unsafe fn match_url_char_32_avx(buf: &[u8]) -> usize {
    // NOTE: This check might be not necessary since this function is only used in
    // `match_uri_vectored` where buffer overflow is taken care of.
    debug_assert!(buf.len() >= 32);

    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;

    // pointer to buffer
    let ptr = buf.as_ptr();

    // %x21-%x7e %x80-%xff
    //
    // Character ranges allowed by this function, can also be interpreted as:
    // 33 =< (x != 127) =< 255
    //
    // Create a vector full of DEL (0x7f) characters.
    let DEL: __m256i = _mm256_set1_epi8(0x7f);
    // Create a vector full of exclamation mark (!) (0x21) characters.
    // Used as lower threshold, characters in URLs cannot be smaller than this.
    let LOW: __m256i = _mm256_set1_epi8(0x21);

    // Load a chunk of 32 bytes from `ptr` as a vector.
    // We can check 32 bytes in parallel at most with AVX2 since
    // YMM registers can only have 256 bits most.
    let dat = _mm256_lddqu_si256(ptr as *const _);

    // unsigned comparison dat >= LOW
    //
    // `_mm256_max_epu8` creates a new vector by comparing vectors `dat` and `LOW`
    // and picks the max. values from each for all indices.
    // So if a byte in `dat` is <= 32, it'll be represented as 33
    // which is the smallest valid character.
    //
    // Then, we compare the new vector with `dat` for equality.
    //
    // `_mm256_cmpeq_epi8` returns a new vector where;
    // * matching bytes are set to 0xFF (all bits set),
    // * nonmatching bytes are set to 0 (no bits set).
    let low = _mm256_cmpeq_epi8(_mm256_max_epu8(dat, LOW), dat);
    // Similar to what we did before, but now invalid characters are set to 0xFF.
    let del = _mm256_cmpeq_epi8(dat, DEL);

    // We glue the both comparisons via `_mm256_andnot_si256`.
    //
    // Since the representation of truthiness differ in these comparisons,
    // we are in need of bitwise NOT to convert valid characters of `del`.
    let bit = _mm256_andnot_si256(del, low);
    // This creates a bitmask from the most significant bit of each byte.
    // Simply, we're converting a vector value to scalar value here.
    let res = _mm256_movemask_epi8(bit) as u32;

    // Count trailing ones to find the first encountered invalid character.
    res.trailing_ones() as usize
}

#[target_feature(enable = "avx2")]
pub unsafe fn match_header_value_vectored(bytes: &mut Bytes) {
    while bytes.as_ref().len() >= 32 {
        let advance = match_header_value_char_32_avx(bytes.as_ref());
        bytes.advance(advance);

        if advance != 32 {
            return;
        }
    }
    // NOTE: use SWAR for <32B, more efficient than falling back to SSE4.2
    super::swar::match_header_value_vectored(bytes);
}

#[inline(always)]
#[allow(non_snake_case)]
#[allow(unused)]
unsafe fn match_header_value_char_32_avx(buf: &[u8]) -> usize {
    debug_assert!(buf.len() >= 32);

    #[cfg(target_arch = "x86")]
    use core::arch::x86::*;
    #[cfg(target_arch = "x86_64")]
    use core::arch::x86_64::*;

    let ptr = buf.as_ptr();

    // %x09 %x20-%x7e %x80-%xff
    // Create a vector full of horizontal tab (\t) (0x09) characters.
    let TAB: __m256i = _mm256_set1_epi8(0x09);
    // Create a vector full of DEL (0x7f) characters.
    let DEL: __m256i = _mm256_set1_epi8(0x7f);
    // Create a vector full of space (0x20) characters.
    let LOW: __m256i = _mm256_set1_epi8(0x20);

    // Load a chunk of 32 bytes from `ptr` as a vector.
    let dat = _mm256_lddqu_si256(ptr as *const _);

    // unsigned comparison dat >= LOW
    //
    // Same as what we do in `match_url_char_32_avx`.
    // This time the lower threshold is set to space character though.
    let low = _mm256_cmpeq_epi8(_mm256_max_epu8(dat, LOW), dat);
    // Check if `dat` includes `TAB` characters.
    let tab = _mm256_cmpeq_epi8(dat, TAB);
    // Check if `dat` includes `DEL` characters.
    let del = _mm256_cmpeq_epi8(dat, DEL);

    // Combine all comparisons together, notice that we're also using OR
    // to connect `low` and `tab` but flip bits of `del`.
    //
    // In the end, this is simply:
    // ~del & (low | tab)
    let bit = _mm256_andnot_si256(del, _mm256_or_si256(low, tab));
    // This creates a bitmask from the most significant bit of each byte.
    // Creates a scalar value from vector value.
    let res = _mm256_movemask_epi8(bit) as u32;

    // Count trailing ones to find the first encountered invalid character.
    res.trailing_ones() as usize
}

#[test]
fn avx2_code_matches_uri_chars_table() {
    if !is_x86_feature_detected!("avx2") {
        return;
    }

    #[allow(clippy::undocumented_unsafe_blocks)]
    unsafe {
        assert!(byte_is_allowed(b'_', match_uri_vectored));

        for (b, allowed) in crate::utils::URI_MAP.iter().copied().enumerate() {
            assert_eq!(
                byte_is_allowed(b as u8, match_uri_vectored),
                allowed,
                "byte_is_allowed({b:?}) should be {allowed:?}"
            );
        }
    }
}

#[test]
fn avx2_code_matches_header_value_chars_table() {
    if !is_x86_feature_detected!("avx2") {
        return;
    }

    #[allow(clippy::undocumented_unsafe_blocks)]
    unsafe {
        assert!(byte_is_allowed(b'_', match_header_value_vectored));

        for (b, allowed) in crate::utils::HEADER_VALUE_MAP.iter().copied().enumerate() {
            assert_eq!(
                byte_is_allowed(b as u8, match_header_value_vectored),
                allowed,
                "byte_is_allowed({b:?}) should be {allowed:?}"
            );
        }
    }
}

#[test]
fn avx2_code_matches_header_name_chars_table() {
    if !is_x86_feature_detected!("avx2") {
        return;
    }

    #[allow(clippy::undocumented_unsafe_blocks)]
    unsafe {
        assert!(byte_is_allowed(b'_', match_header_name_vectored));

        for (b, allowed) in crate::utils::TOKEN_MAP.iter().copied().enumerate() {
            assert_eq!(
                byte_is_allowed(b as u8, match_header_name_vectored),
                allowed,
                "byte_is_allowed({b:?}) should be {allowed:?}"
            );
            // the 16-byte step
            assert_eq!(
                byte_is_allowed_16(b as u8, match_header_name_vectored),
                allowed,
                "byte_is_allowed_16({b:?}) should be {allowed:?}"
            );
        }
    }
}

#[cfg(test)]
unsafe fn byte_is_allowed_16(byte: u8, f: unsafe fn(bytes: &mut Bytes<'_, '_>)) -> bool {
    let mut st = crate::State::default();
    let mut slice = [b'_'; 20];
    slice[10] = byte;
    let mut bytes = Bytes::new(&slice, &mut st);

    f(&mut bytes);

    match bytes.cursor() - bytes.start() {
        20 => true,
        10 => false,
        _ => unreachable!(),
    }
}

#[cfg(test)]
unsafe fn byte_is_allowed(byte: u8, f: unsafe fn(bytes: &mut Bytes<'_, '_>)) -> bool {
    let mut st = crate::State::default();
    let slice = [
        b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_',
        b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_', b'_', byte, b'_', b'_', b'_',
        b'_', b'_',
    ];
    let mut bytes = Bytes::new(&slice, &mut st);

    f(&mut bytes);

    match bytes.cursor() - bytes.start() {
        32 => true,
        26 => false,
        _ => unreachable!(),
    }
}