fastsearch_simd.patch

diff --git a/Objects/stringlib/fastsearch.h b/Objects/stringlib/fastsearch.h
index 26bb055..1c96ecd 100644
--- a/Objects/stringlib/fastsearch.h
+++ b/Objects/stringlib/fastsearch.h
@@ -769,6 +769,154 @@ STRINGLIB(count_char_no_maxcount)(const STRINGLIB_CHAR *s, Py_ssize_t n,
 }
 
 
+/* SIMD-accelerated forward search / count for byte-width strings.
+
+   Uses the classic "first byte + last byte" filter (as in glibc memmem, ripgrep
+   memchr::memmem and StringZilla): a match can only start where both p[0] and
+   p[m-1] line up, so one vector instruction rejects 16 candidate start positions
+   at once and only the survivors are verified with a short compare. To preserve
+   CPython's O(n) worst-case guarantee against adversarial needles (e.g. "aaab"
+   in "aaaa..."), it falls back to the two-way algorithm once verification work
+   grows disproportionate, mirroring adaptive_find.
+
+   STRINGLIB_SIMD_MASK16(h, i) yields a 16-bit mask whose bit k is set iff the
+   first+last filter passes at position i+k. Only the arch-specific mask differs;
+   all match/verify/count/fallback logic below is shared. */
+/* SSE2 is baseline on every x86-64 target (GCC/Clang define __SSE2__; MSVC does
+   not, but guarantees it under _M_X64) and NEON is baseline on aarch64, so a
+   compile-time guard needs no runtime CPU detection for these two. Wider ISAs
+   (AVX2, SVE) would; see python/cpython#125022. Other compilers/targets and the
+   wide (UCS2/UCS4) str representations fall through to the scalar path unchanged. */
+#if STRINGLIB_SIZEOF_CHAR == 1 && \
+    (defined(__ARM_NEON) || defined(__SSE2__) || (defined(_MSC_VER) && defined(_M_X64)))
+#define STRINGLIB_HAVE_SIMD_FIND 1  /* local marker, undef below */
+#if defined(__ARM_NEON)
+#include <arm_neon.h>
+static inline uint32_t
+STRINGLIB(_simd_mask16)(const uint8_t *h, Py_ssize_t i, Py_ssize_t last,
+                        uint8x16_t vf, uint8x16_t vl)
+{
+    uint8x16_t eq = vandq_u8(vceqq_u8(vld1q_u8(h + i), vf),
+                             vceqq_u8(vld1q_u8(h + i + last), vl));
+    /* NEON movemask: AND each all-ones lane with a per-lane bit, then a pair of
+       horizontal adds collapses the two halves into a 16-bit mask. This mirrors
+       SSE2's _mm_movemask_epi8 in a few instructions, with no per-lane branch. */
+    const uint8x16_t bits = {1, 2, 4, 8, 16, 32, 64, 128,
+                             1, 2, 4, 8, 16, 32, 64, 128};
+    uint8x16_t m = vandq_u8(eq, bits);
+    return (uint32_t)vaddv_u8(vget_low_u8(m)) |
+           ((uint32_t)vaddv_u8(vget_high_u8(m)) << 8);
+}
+#else  /* SSE2: GCC/Clang via __SSE2__, or MSVC x64 via _M_X64 */
+#if defined(_MSC_VER)
+#include <intrin.h>      /* SSE2 intrinsics + _BitScanForward */
+#else
+#include <emmintrin.h>
+#endif
+static inline uint32_t
+STRINGLIB(_simd_mask16)(const uint8_t *h, Py_ssize_t i, Py_ssize_t last,
+                        __m128i vf, __m128i vl)
+{
+    __m128i eq = _mm_and_si128(
+        _mm_cmpeq_epi8(_mm_loadu_si128((const __m128i *)(h + i)), vf),
+        _mm_cmpeq_epi8(_mm_loadu_si128((const __m128i *)(h + i + last)), vl));
+    return (uint32_t)_mm_movemask_epi8(eq);
+}
+#endif
+
+/* Index of the lowest set bit; mask is always non-zero at the call site. */
+static inline int
+STRINGLIB(_simd_ctz)(uint32_t mask)
+{
+#if defined(_MSC_VER)
+    unsigned long i;
+    _BitScanForward(&i, mask);
+    return (int)i;
+#else
+    return __builtin_ctz(mask);
+#endif
+}
+
+static Py_ssize_t
+STRINGLIB(_simd_find)(const STRINGLIB_CHAR* s, Py_ssize_t n,
+                      const STRINGLIB_CHAR* p, Py_ssize_t m,
+                      Py_ssize_t maxcount, int mode)
+{
+    const uint8_t *h = (const uint8_t *)s;
+    const uint8_t *pat = (const uint8_t *)p;
+    const Py_ssize_t last = m - 1;
+    const Py_ssize_t end = n - m + 1;          /* number of valid start positions */
+#if defined(__ARM_NEON)
+    const uint8x16_t vf = vdupq_n_u8(pat[0]), vl = vdupq_n_u8(pat[last]);
+#else
+    const __m128i vf = _mm_set1_epi8((char)pat[0]), vl = _mm_set1_epi8((char)pat[last]);
+#endif
+    Py_ssize_t count = 0, hits = 0, next = 0, i = 0;
+
+    while (i + 16 <= end) {
+        uint32_t mask = STRINGLIB(_simd_mask16)(h, i, last, vf, vl);
+        while (mask) {
+            int k = STRINGLIB(_simd_ctz)(mask);            /* index of lowest set bit */
+            mask &= mask - 1;
+            Py_ssize_t pos = i + k;
+            Py_ssize_t j = 1;
+            while (j < last && h[pos + j] == pat[j]) {
+                j++;
+            }
+            hits += j;
+            if (j == last) {                    /* full match at pos */
+                if (mode != FAST_COUNT) {
+                    return pos;
+                }
+                if (pos >= next) {              /* non-overlapping count */
+                    count++;
+                    if (count == maxcount) {
+                        return maxcount;
+                    }
+                    next = pos + m;
+                }
+            }
+            /* adaptive fallback keeps the linear worst-case bound */
+            if (hits > m + (i >> 2) && end - pos > 2000) {
+                if (mode == FAST_SEARCH) {
+                    Py_ssize_t r = STRINGLIB(_two_way_find)(s + pos, n - pos, p, m);
+                    return r == -1 ? -1 : r + pos;
+                }
+                Py_ssize_t startp = pos > next ? pos : next;
+                Py_ssize_t r = STRINGLIB(_two_way_count)(s + startp, n - startp, p, m,
+                                                         maxcount - count);
+                return r + count;
+            }
+        }
+        i += 16;
+    }
+    /* scalar tail for the final < 16 candidate positions */
+    for (; i < end; i++) {
+        if (h[i] != pat[0] || h[i + last] != pat[last]) {
+            continue;
+        }
+        Py_ssize_t j = 1;
+        while (j < last && h[i + j] == pat[j]) {
+            j++;
+        }
+        if (j == last) {
+            if (mode != FAST_COUNT) {
+                return i;
+            }
+            if (i >= next) {
+                count++;
+                if (count == maxcount) {
+                    return maxcount;
+                }
+                next = i + m;
+            }
+        }
+    }
+    return mode == FAST_COUNT ? count : -1;
+}
+#endif  /* STRINGLIB_SIZEOF_CHAR == 1 && (NEON || SSE2) */
+
+
 Py_LOCAL_INLINE(Py_ssize_t)
 FASTSEARCH(const STRINGLIB_CHAR* s, Py_ssize_t n,
            const STRINGLIB_CHAR* p, Py_ssize_t m,
@@ -797,6 +945,14 @@ FASTSEARCH(const STRINGLIB_CHAR* s, Py_ssize_t n,
     }
 
     if (mode != FAST_RSEARCH) {
+#if STRINGLIB_SIZEOF_CHAR == 1 && \
+    (defined(__ARM_NEON) || defined(__SSE2__) || (defined(_MSC_VER) && defined(_M_X64)))
+        /* SIMD forward search/count wins once there is enough haystack to amortise
+           the vector setup; below that the scalar Bloom/Horspool path is fine. */
+        if (m >= 2 && n >= 64) {
+            return STRINGLIB(_simd_find)(s, n, p, m, maxcount, mode);
+        }
+#endif
         if (n < 2500 || (m < 100 && n < 30000) || m < 6) {
             return STRINGLIB(default_find)(s, n, p, m, maxcount, mode);
         }