    1
    2
    3
    4
    5
    6
    7
    8
    9
    10
    11
    12
    13
    14
    15
    16
    17
    18
    19
    20
    21
    22
    23
    24
    25
    26
    27
    28
    29
    30
    31
    32
    33
    34
    35
    36
    37
    38
    39
    40
    41
    42
    43
    44
    45
    46
    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    86
    87
    88
    89
    90
    91
    92
    93
    94
    95
    96
    97
    98
    99
    100
    101
    102
    103
    104
    105
    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    206
    207

    use core::mem::size_of;

    use crate::memmem::{
        prefilter::{PrefilterFnTy, PrefilterState},
        vector::Vector,
        NeedleInfo,
    };

    /// The implementation of the forward vector accelerated candidate finder.
    ///
    /// This is inspired by the "generic SIMD" algorithm described here:
    /// http://0x80.pl/articles/simd-strfind.html#algorithm-1-generic-simd
    ///
    /// The main difference is that this is just a prefilter. That is, it reports
    /// candidates once they are seen and doesn't attempt to confirm them. Also,
    /// the bytes this routine uses to check for candidates are selected based on
    /// an a priori background frequency distribution. This means that on most
    /// haystacks, this will on average spend more time in vectorized code than you
    /// would if you just selected the first and last bytes of the needle.
    ///
    /// Note that a non-prefilter variant of this algorithm can be found in the
    /// parent module, but it only works on smaller needles.
    ///
    /// `prestate`, `ninfo`, `haystack` and `needle` are the four prefilter
    /// function parameters. `fallback` is a prefilter that is used if the haystack
    /// is too small to be handled with the given vector size.
    ///
    /// This routine is not safe because it is intended for callers to specialize
    /// this with a particular vector (e.g., __m256i) and then call it with the
    /// relevant target feature (e.g., avx2) enabled.
    ///
    /// # Panics
    ///
    /// If `needle.len() <= 1`, then this panics.
    ///
    /// # Safety
    ///
    /// Since this is meant to be used with vector functions, callers need to
    /// specialize this inside of a function with a `target_feature` attribute.
    /// Therefore, callers must ensure that whatever target feature is being used
    /// supports the vector functions that this function is specialized for. (For
    /// the specific vector functions used, see the Vector trait implementations.)
    #[inline(always)]
    pub(crate) unsafe fn find<V: Vector>(
        prestate: &mut PrefilterState,
        ninfo: &NeedleInfo,
        haystack: &[u8],
        needle: &[u8],
        fallback: PrefilterFnTy,
    ) -> Option<usize> {
        assert!(needle.len() >= 2, "needle must be at least 2 bytes");
        let (rare1i, rare2i) = ninfo.rarebytes.as_rare_ordered_usize();
        let min_haystack_len = rare2i + size_of::<V>();
        if haystack.len() < min_haystack_len {
            return fallback(prestate, ninfo, haystack, needle);
        }

        let start_ptr = haystack.as_ptr();
        let end_ptr = start_ptr.add(haystack.len());
        let max_ptr = end_ptr.sub(min_haystack_len);
        let mut ptr = start_ptr;

        let rare1chunk = V::splat(needle[rare1i]);
        let rare2chunk = V::splat(needle[rare2i]);

        // N.B. I did experiment with unrolling the loop to deal with size(V)
        // bytes at a time and 2*size(V) bytes at a time. The double unroll
        // was marginally faster while the quadruple unroll was unambiguously
        // slower. In the end, I decided the complexity from unrolling wasn't
        // worth it. I used the memmem/krate/prebuilt/huge-en/ benchmarks to
        // compare.
        while ptr <= max_ptr {
            let m = find_in_chunk2(ptr, rare1i, rare2i, rare1chunk, rare2chunk);
            if let Some(chunki) = m {
                return Some(matched(prestate, start_ptr, ptr, chunki));
            }
            ptr = ptr.add(size_of::<V>());
        }
        if ptr < end_ptr {
            // This routine immediately quits if a candidate match is found.
            // That means that if we're here, no candidate matches have been
            // found at or before 'ptr'. Thus, we don't need to mask anything
            // out even though we might technically search part of the haystack
            // that we've already searched (because we know it can't match).
            ptr = max_ptr;
            let m = find_in_chunk2(ptr, rare1i, rare2i, rare1chunk, rare2chunk);
            if let Some(chunki) = m {
                return Some(matched(prestate, start_ptr, ptr, chunki));
            }
        }
        prestate.update(haystack.len());
        None
    }

    // Below are two different techniques for checking whether a candidate
    // match exists in a given chunk or not. find_in_chunk2 checks two bytes
    // where as find_in_chunk3 checks three bytes. The idea behind checking
    // three bytes is that while we do a bit more work per iteration, we
    // decrease the chances of a false positive match being reported and thus
    // make the search faster overall. This actually works out for the
    // memmem/krate/prebuilt/huge-en/never-all-common-bytes benchmark, where
    // using find_in_chunk3 is about 25% faster than find_in_chunk2. However,
    // it turns out that find_in_chunk2 is faster for all other benchmarks, so
    // perhaps the extra check isn't worth it in practice.
    //
    // For now, we go with find_in_chunk2, but we leave find_in_chunk3 around
    // to make it easy to switch to and benchmark when possible.

    /// Search for an occurrence of two rare bytes from the needle in the current
    /// chunk pointed to by ptr.
    ///
    /// rare1chunk and rare2chunk correspond to vectors with the rare1 and rare2
    /// bytes repeated in each 8-bit lane, respectively.
    ///
    /// # Safety
    ///
    /// It must be safe to do an unaligned read of size(V) bytes starting at both
    /// (ptr + rare1i) and (ptr + rare2i).
    #[inline(always)]
    unsafe fn find_in_chunk2<V: Vector>(
        ptr: *const u8,
        rare1i: usize,
        rare2i: usize,
        rare1chunk: V,
        rare2chunk: V,
    ) -> Option<usize> {
        let chunk0 = V::load_unaligned(ptr.add(rare1i));
        let chunk1 = V::load_unaligned(ptr.add(rare2i));

        let eq0 = chunk0.cmpeq(rare1chunk);
        let eq1 = chunk1.cmpeq(rare2chunk);

        let match_offsets = eq0.and(eq1).movemask();
        if match_offsets == 0 {
            return None;
        }
        Some(match_offsets.trailing_zeros() as usize)
    }

    /// Search for an occurrence of two rare bytes and the first byte (even if one
    /// of the rare bytes is equivalent to the first byte) from the needle in the
    /// current chunk pointed to by ptr.
    ///
    /// firstchunk, rare1chunk and rare2chunk correspond to vectors with the first,
    /// rare1 and rare2 bytes repeated in each 8-bit lane, respectively.
    ///
    /// # Safety
    ///
    /// It must be safe to do an unaligned read of size(V) bytes starting at ptr,
    /// (ptr + rare1i) and (ptr + rare2i).
    #[allow(dead_code)]
    #[inline(always)]
    unsafe fn find_in_chunk3<V: Vector>(
        ptr: *const u8,
        rare1i: usize,
        rare2i: usize,
        firstchunk: V,
        rare1chunk: V,
        rare2chunk: V,
    ) -> Option<usize> {
        let chunk0 = V::load_unaligned(ptr);
        let chunk1 = V::load_unaligned(ptr.add(rare1i));
        let chunk2 = V::load_unaligned(ptr.add(rare2i));

        let eq0 = chunk0.cmpeq(firstchunk);
        let eq1 = chunk1.cmpeq(rare1chunk);
        let eq2 = chunk2.cmpeq(rare2chunk);

        let match_offsets = eq0.and(eq1).and(eq2).movemask();
        if match_offsets == 0 {
            return None;
        }
        Some(match_offsets.trailing_zeros() as usize)
    }

    /// Accepts a chunk-relative offset and returns a haystack relative offset
    /// after updating the prefilter state.
    ///
    /// Why do we use this unlineable function when a search completes? Well,
    /// I don't know. Really. Obviously this function was not here initially.
    /// When doing profiling, the codegen for the inner loop here looked bad and
    /// I didn't know why. There were a couple extra 'add' instructions and an
    /// extra 'lea' instruction that I couldn't explain. I hypothesized that the
    /// optimizer was having trouble untangling the hot code in the loop from the
    /// code that deals with a candidate match. By putting the latter into an
    /// unlineable function, it kind of forces the issue and it had the intended
    /// effect: codegen improved measurably. It's good for a ~10% improvement
    /// across the board on the memmem/krate/prebuilt/huge-en/ benchmarks.
    #[cold]
    #[inline(never)]
    fn matched(
        prestate: &mut PrefilterState,
        start_ptr: *const u8,
        ptr: *const u8,
        chunki: usize,
    ) -> usize {
        let found = diff(ptr, start_ptr) + chunki;
        prestate.update(found);
        found
    }

    /// Subtract `b` from `a` and return the difference. `a` must be greater than
    /// or equal to `b`.
    fn diff(a: *const u8, b: *const u8) -> usize {
        debug_assert!(a >= b);
        (a as usize) - (b as usize)
    }