[mypyc] feat: optimize C code for str.count

Author: BobTheBuidler

mypyc/lib-rt/str_ops.c

@@ -532,28 +532,13 @@ PyObject *CPy_Encode(PyObject *obj, PyObject *encoding, PyObject *errors) {
     }
 }
 
-Py_ssize_t CPyStr_Count(PyObject *unicode, PyObject *substring, CPyTagged start) {
-    Py_ssize_t temp_start = CPyTagged_AsSsize_t(start);
-    if (temp_start == -1 && PyErr_Occurred()) {
-        PyErr_SetString(PyExc_OverflowError, CPYTHON_LARGE_INT_ERRMSG);
-        return -1;
-    }
+Py_ssize_t CPyStr_Count(PyObject *unicode, PyObject *substring, Py_ssize_t start) {
     Py_ssize_t end = PyUnicode_GET_LENGTH(unicode);
-    return PyUnicode_Count(unicode, substring, temp_start, end);
+    return PyUnicode_Count(unicode, substring, start, end);
 }
 
-Py_ssize_t CPyStr_CountFull(PyObject *unicode, PyObject *substring, CPyTagged start, CPyTagged end) {
-    Py_ssize_t temp_start = CPyTagged_AsSsize_t(start);
-    if (temp_start == -1 && PyErr_Occurred()) {
-        PyErr_SetString(PyExc_OverflowError, CPYTHON_LARGE_INT_ERRMSG);
-        return -1;
-    }
-    Py_ssize_t temp_end = CPyTagged_AsSsize_t(end);
-    if (temp_end == -1 && PyErr_Occurred()) {
-        PyErr_SetString(PyExc_OverflowError, CPYTHON_LARGE_INT_ERRMSG);
-        return -1;
-    }
-    return PyUnicode_Count(unicode, substring, temp_start, temp_end);
+Py_ssize_t CPyStr_CountFull(PyObject *unicode, PyObject *substring, Py_ssize_t start, Py_ssize_t end) {
+    return PyUnicode_Count(unicode, substring, start, end);
 }
 
 

mypyc/primitives/str_ops.py

@@ -318,7 +318,7 @@
 # str.count(substring, start)
 method_op(
     name="count",
-    arg_types=[str_rprimitive, str_rprimitive, int_rprimitive],
+    arg_types=[str_rprimitive, str_rprimitive, c_pyssize_t_rprimitive],
     return_type=c_pyssize_t_rprimitive,
     c_function_name="CPyStr_Count",
     error_kind=ERR_NEG_INT,
@@ -327,7 +327,7 @@
 # str.count(substring, start, end)
 method_op(
     name="count",
-    arg_types=[str_rprimitive, str_rprimitive, int_rprimitive, int_rprimitive],
+    arg_types=[str_rprimitive, str_rprimitive, c_pyssize_t_rprimitive, c_pyssize_t_rprimitive],
     return_type=c_pyssize_t_rprimitive,
     c_function_name="CPyStr_CountFull",
     error_kind=ERR_NEG_INT,