<commit_before>use crate::postings::compression::AlignedBuffer;

/// This modules define the logic used to search for a doc in a given
/// block. (at most 128 docs)
///
/// Searching within a block is a hotspot when running intersection.
/// so it was worth defining it in its own module.

#[cfg(target_arch = "x86_64")]
mod sse2 {
    use crate::postings::compression::{AlignedBuffer, COMPRESSION_BLOCK_SIZE};
    use std::arch::x86_64::__m128i as DataType;
    use std::arch::x86_64::_mm_add_epi32 as op_add;
    use std::arch::x86_64::_mm_cmplt_epi32 as op_lt;
    use std::arch::x86_64::_mm_load_si128 as op_load; // requires 128-bits alignment
    use std::arch::x86_64::_mm_set1_epi32 as set1;
    use std::arch::x86_64::_mm_setzero_si128 as set0;
    use std::arch::x86_64::_mm_sub_epi32 as op_sub;
    use std::arch::x86_64::{_mm_cvtsi128_si32, _mm_shuffle_epi32};

    const MASK1: i32 = 78;
    const MASK2: i32 = 177;

    /// Performs an exhaustive linear search over the
    ///
    /// There is no early exit here. We simply count the
    /// number of elements that are `< target`.
    pub(crate) fn linear_search_sse2_128(arr: &AlignedBuffer, target: u32) -> usize {
        unsafe {
            let ptr = arr as *const AlignedBuffer as *const DataType;
            let vkey = set1(target as i32);
            let mut cnt = set0();
            // We work over 4 `__m128i` at a time.
            // A single `__m128i` actual contains 4 `u32`.
            for i in 0..(COMPRESSION_BLOCK_SIZE as isize) / (4 * 4) {
                let cmp1 = op_lt(op_load(ptr.offset(i * 4)), vkey);
                let cmp2 = op_lt(op_load(ptr.offset(i * 4 + 1)), vkey);
                let cmp3 = op_lt(op_load(ptr.offset(i * 4 + 2)), vkey);
                let cmp4 = op_lt(op_load(ptr.offset(i * 4 + 3)), vkey);
                let sum = op_add(op_add(cmp1, cmp2), op_add(cmp3, cmp4));
                cnt = op_sub(cnt, sum);
            }
            cnt = op_add(cnt, _mm_shuffle_epi32(cnt, MASK1));
            cnt = op_add(cnt, _mm_shuffle_epi32(cnt, MASK2));
            _mm_cvtsi128_si32(cnt) as usize
        }
    }

    #[cfg(test)]
    mod test {
        use super::linear_search_sse2_128;
        use crate::postings::compression::{AlignedBuffer, COMPRESSION_BLOCK_SIZE};

        #[test]
        fn test_linear_search_sse2_128_u32() {
            let mut block = [0u32; COMPRESSION_BLOCK_SIZE];
            for el in 0u32..128u32 {
                block[el as usize] = el * 2 + 1 << 18;
            }
            let target = block[64] + 1;
            assert_eq!(linear_search_sse2_128(&AlignedBuffer(block), target), 65);
        }
    }
}

/// This `linear search` browser exhaustively through the array.
/// but the early exit is very difficult to predict.
///
/// Coupled with `exponential search` this function is likely
/// to be called with the same `len`
fn linear_search(arr: &[u32], target: u32) -> usize {
    arr.iter().map(|&el| if el < target { 1 } else { 0 }).sum()
}

fn exponential_search(arr: &[u32], target: u32) -> (usize, usize) {
    let end = arr.len();
    let mut begin = 0;
    for &pivot in &[1, 3, 7, 15, 31, 63] {
        if pivot >= end {
            break;
        }
        if arr[pivot] > target {
            return (begin, pivot);
        }
        begin = pivot;
    }
    (begin, end)
}

fn galloping(block_docs: &[u32], target: u32) -> usize {
    let (start, end) = exponential_search(&block_docs, target);
    start + linear_search(&block_docs[start..end], target)
}

/// Tantivy may rely on SIMD instructions to search for a specific document within
/// a given block.
#[derive(Clone, Copy, PartialEq)]
pub enum BlockSearcher {
    #[cfg(target_arch = "x86_64")]
    SSE2,
    Scalar,
}

impl BlockSearcher {
    /// Search the first index containing an element greater or equal to
    /// the target.
    ///
    /// The results should be equivalent to
    /// ```ignore
    /// block[..]
    //       .iter()
    //       .take_while(|&&val| val < target)
    //       .count()
    /// ```
    ///
    /// The `start` argument is just used to hint that the response is
    /// greater than beyond `start`. The implementation may or may not use
    /// it for optimization.
    ///
    /// # Assumption
    ///
    /// The array len is > start.
    /// The block is sorted
    /// The target is assumed greater or equal to the `arr[start]`.
    /// The target is assumed smaller or equal to the last element of the block.
    ///
    /// Currently the scalar implementation starts by an exponential search, and
    /// then operates a linear search in the result subarray.<commit_msg>ignore -> compile_fail<commit_after>use crate::postings::compression::AlignedBuffer;

/// This modules define the logic used to search for a doc in a given
/// block. (at most 128 docs)
///
/// Searching within a block is a hotspot when running intersection.
/// so it was worth defining it in its own module.

#[cfg(target_arch = "x86_64")]
mod sse2 {
    use crate::postings::compression::{AlignedBuffer, COMPRESSION_BLOCK_SIZE};
    use std::arch::x86_64::__m128i as DataType;
    use std::arch::x86_64::_mm_add_epi32 as op_add;
    use std::arch::x86_64::_mm_cmplt_epi32 as op_lt;
    use std::arch::x86_64::_mm_load_si128 as op_load; // requires 128-bits alignment
    use std::arch::x86_64::_mm_set1_epi32 as set1;
    use std::arch::x86_64::_mm_setzero_si128 as set0;
    use std::arch::x86_64::_mm_sub_epi32 as op_sub;
    use std::arch::x86_64::{_mm_cvtsi128_si32, _mm_shuffle_epi32};

    const MASK1: i32 = 78;
    const MASK2: i32 = 177;

    /// Performs an exhaustive linear search over the
    ///
    /// There is no early exit here. We simply count the
    /// number of elements that are `< target`.
    pub(crate) fn linear_search_sse2_128(arr: &AlignedBuffer, target: u32) -> usize {
        unsafe {
            let ptr = arr as *const AlignedBuffer as *const DataType;
            let vkey = set1(target as i32);
            let mut cnt = set0();
            // We work over 4 `__m128i` at a time.
            // A single `__m128i` actual contains 4 `u32`.
            for i in 0..(COMPRESSION_BLOCK_SIZE as isize) / (4 * 4) {
                let cmp1 = op_lt(op_load(ptr.offset(i * 4)), vkey);
                let cmp2 = op_lt(op_load(ptr.offset(i * 4 + 1)), vkey);
                let cmp3 = op_lt(op_load(ptr.offset(i * 4 + 2)), vkey);
                let cmp4 = op_lt(op_load(ptr.offset(i * 4 + 3)), vkey);
                let sum = op_add(op_add(cmp1, cmp2), op_add(cmp3, cmp4));
                cnt = op_sub(cnt, sum);
            }
            cnt = op_add(cnt, _mm_shuffle_epi32(cnt, MASK1));
            cnt = op_add(cnt, _mm_shuffle_epi32(cnt, MASK2));
            _mm_cvtsi128_si32(cnt) as usize
        }
    }

    #[cfg(test)]
    mod test {
        use super::linear_search_sse2_128;
        use crate::postings::compression::{AlignedBuffer, COMPRESSION_BLOCK_SIZE};

        #[test]
        fn test_linear_search_sse2_128_u32() {
            let mut block = [0u32; COMPRESSION_BLOCK_SIZE];
            for el in 0u32..128u32 {
                block[el as usize] = el * 2 + 1 << 18;
            }
            let target = block[64] + 1;
            assert_eq!(linear_search_sse2_128(&AlignedBuffer(block), target), 65);
        }
    }
}

/// This `linear search` browser exhaustively through the array.
/// but the early exit is very difficult to predict.
///
/// Coupled with `exponential search` this function is likely
/// to be called with the same `len`
fn linear_search(arr: &[u32], target: u32) -> usize {
    arr.iter().map(|&el| if el < target { 1 } else { 0 }).sum()
}

fn exponential_search(arr: &[u32], target: u32) -> (usize, usize) {
    let end = arr.len();
    let mut begin = 0;
    for &pivot in &[1, 3, 7, 15, 31, 63] {
        if pivot >= end {
            break;
        }
        if arr[pivot] > target {
            return (begin, pivot);
        }
        begin = pivot;
    }
    (begin, end)
}

fn galloping(block_docs: &[u32], target: u32) -> usize {
    let (start, end) = exponential_search(&block_docs, target);
    start + linear_search(&block_docs[start..end], target)
}

/// Tantivy may rely on SIMD instructions to search for a specific document within
/// a given block.
#[derive(Clone, Copy, PartialEq)]
pub enum BlockSearcher {
    #[cfg(target_arch = "x86_64")]
    SSE2,
    Scalar,
}

impl BlockSearcher {
    /// Search the first index containing an element greater or equal to
    /// the target.
    ///
    /// The results should be equivalent to
    /// ```compile_fail
    /// block[..]
    //       .iter()
    //       .take_while(|&&val| val < target)
    //       .count()
    /// ```
    ///
    /// The `start` argument is just used to hint that the response is
    /// greater than beyond `start`. The implementation may or may not use
    /// it for optimization.
    ///
    /// # Assumption
    ///
    /// The array len is > start.
    /// The block is sorted
    /// The target is assumed greater or equal to the `arr[start]`.
    /// The target is assumed smaller or equal to the last element of the block.
    ///
    /// Currently the scalar implementation starts by an exponential search, and
    /// then operates a linear search in the result subarray.
{'fraction_non_alphanumeric': 0.12848380501452705, 'fraction_numerical': 0.04153664048208329, 'mean_word_length': 2.2748414376321353, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 18, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '492195', 'n_tokens_mistral': 3286, 'n_tokens_neox': 3059, 'n_words': 938}