gh-120397: Optimize str.count() for single characters (#120398)

1 files changed, 19 insertions, 0 deletions

diff --git a/Objects/stringlib/fastsearch.h b/Objects/stringlib/fastsearch.h
index 309ed15..05e700b 100644
--- a/Objects/stringlib/fastsearch.h
+++ b/Objects/stringlib/fastsearch.h
@@ -753,6 +753,22 @@ STRINGLIB(count_char)(const STRINGLIB_CHAR *s, Py_ssize_t n,
 }
 
 
+static inline Py_ssize_t
+STRINGLIB(count_char_no_maxcount)(const STRINGLIB_CHAR *s, Py_ssize_t n,
+                                  const STRINGLIB_CHAR p0)
+/* A specialized function of count_char that does not cut off at a maximum.
+   As a result, the compiler is able to vectorize the loop. */
+{
+    Py_ssize_t count = 0;
+    for (Py_ssize_t i = 0; i < n; i++) {
+        if (s[i] == p0) {
+            count++;
+        }
+    }
+    return count;
+}
+
+
 Py_LOCAL_INLINE(Py_ssize_t)
 FASTSEARCH(const STRINGLIB_CHAR* s, Py_ssize_t n,
            const STRINGLIB_CHAR* p, Py_ssize_t m,
@@ -773,6 +789,9 @@ FASTSEARCH(const STRINGLIB_CHAR* s, Py_ssize_t n,
         else if (mode == FAST_RSEARCH)
             return STRINGLIB(rfind_char)(s, n, p[0]);
         else {
+            if (maxcount == PY_SSIZE_T_MAX) {
+                return STRINGLIB(count_char_no_maxcount)(s, n, p[0]);
+            }
             return STRINGLIB(count_char)(s, n, p[0], maxcount);
         }
     }