[mypyc] optimise startswith and endswith (#9557)

Relates to mypyc/mypyc#644.

Author: r3m0t

mypyc/lib-rt/str_ops.c

@@ -53,6 +53,18 @@ PyObject *CPyStr_Split(PyObject *str, PyObject *sep, CPyTagged max_split)
     return PyUnicode_Split(str, sep, temp_max_split);
 }
 
+bool CPyStr_Startswith(PyObject *self, PyObject *subobj) {
+    Py_ssize_t start = 0;
+    Py_ssize_t end = PyUnicode_GET_LENGTH(self);
+    return PyUnicode_Tailmatch(self, subobj, start, end, -1);
+}
+
+bool CPyStr_Endswith(PyObject *self, PyObject *subobj) {
+    Py_ssize_t start = 0;
+    Py_ssize_t end = PyUnicode_GET_LENGTH(self);
+    return PyUnicode_Tailmatch(self, subobj, start, end, 1);
+}
+
 /* This does a dodgy attempt to append in place  */
 PyObject *CPyStr_Append(PyObject *o1, PyObject *o2) {
     PyUnicode_Append(&o1, o2);

mypyc/primitives/str_ops.py

@@ -5,7 +5,7 @@
 from mypyc.ir.ops import ERR_MAGIC, ERR_NEVER
 from mypyc.ir.rtypes import (
     RType, object_rprimitive, str_rprimitive, int_rprimitive, list_rprimitive,
-    c_int_rprimitive, pointer_rprimitive
+    c_int_rprimitive, pointer_rprimitive, bool_rprimitive
 )
 from mypyc.primitives.registry import (
     c_method_op, c_binary_op, c_function_op,
@@ -43,6 +43,24 @@
     error_kind=ERR_MAGIC
 )
 
+# str.startswith(str)
+c_method_op(
+    name='startswith',
+    arg_types=[str_rprimitive, str_rprimitive],
+    return_type=bool_rprimitive,
+    c_function_name='CPyStr_Startswith',
+    error_kind=ERR_NEVER
+)
+
+# str.endswith(str)
+c_method_op(
+    name='endswith',
+    arg_types=[str_rprimitive, str_rprimitive],
+    return_type=bool_rprimitive,
+    c_function_name='CPyStr_Endswith',
+    error_kind=ERR_NEVER
+)
+
 # str[index] (for an int index)
 c_method_op(
     name='__getitem__',

mypyc/test-data/fixtures/ir.py

@@ -73,6 +73,8 @@ def strip (self, item: str) -> str: pass
     def join(self, x: Iterable[str]) -> str: pass
     def format(self, *args: Any, **kwargs: Any) -> str: ...
     def upper(self) -> str: pass
+    def startswith(self, x: str, start: int=..., end: int=...) -> bool: pass
+    def endswith(self, x: str, start: int=..., end: int=...) -> bool: pass
 
 class float:
     def __init__(self, x: object) -> None: pass

mypyc/test-data/run-strings.test

@@ -1,6 +1,7 @@
 # Test cases for strings (compile and run)
 
 [case testStr]
+from typing import Tuple
 def f() -> str:
     return 'some string'
 def g() -> str:
@@ -17,9 +18,11 @@ def eq(x: str) -> int:
     elif x != 'bar':
         return 1
     return 2
+def match(x: str, y: str) -> Tuple[bool, bool]:
+    return (x.startswith(y), x.endswith(y))
 
 [file driver.py]
-from native import f, g, tostr, booltostr, concat, eq
+from native import f, g, tostr, booltostr, concat, eq, match
 assert f() == 'some string'
 assert g() == 'some\a \v \t \x7f " \n \0string 🐍'
 assert tostr(57) == '57'
@@ -32,6 +35,11 @@ assert eq('bar') == 2
 
 assert int(tostr(0)) == 0
 assert int(tostr(20)) == 20
+assert match('', '') == (True, True)
+assert match('abc', '') == (True, True)
+assert match('abc', 'a') == (True, False)
+assert match('abc', 'c') == (False, True)
+assert match('', 'abc') == (False, False)
 
 [case testStringOps]
 from typing import List, Optional