From 52af4726e4f22d15bd730f2ed97aa79ffbbaab41 Mon Sep 17 00:00:00 2001
From: Tomer Chachamu <tomer.chachamu@gmail.com>
Date: Thu, 8 Oct 2020 14:21:26 +0100
Subject: [PATCH] [mypyc] optimise startswith and endswith

---
 mypyc/lib-rt/str_ops.c           | 12 ++++++++++++
 mypyc/primitives/str_ops.py      | 20 +++++++++++++++++++-
 mypyc/test-data/fixtures/ir.py   |  2 ++
 mypyc/test-data/run-strings.test | 10 +++++++++-
 4 files changed, 42 insertions(+), 2 deletions(-)

diff --git a/mypyc/lib-rt/str_ops.c b/mypyc/lib-rt/str_ops.c
index fe892bb110b60..87e473e275745 100644
--- a/mypyc/lib-rt/str_ops.c
+++ b/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);
diff --git a/mypyc/primitives/str_ops.py b/mypyc/primitives/str_ops.py
index 51b1056cdca2a..b0261a9b4d983 100644
--- a/mypyc/primitives/str_ops.py
+++ b/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__',
diff --git a/mypyc/test-data/fixtures/ir.py b/mypyc/test-data/fixtures/ir.py
index 66d1c58137437..4ffefb7432de8 100644
--- a/mypyc/test-data/fixtures/ir.py
+++ b/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
diff --git a/mypyc/test-data/run-strings.test b/mypyc/test-data/run-strings.test
index 50960aeac1c4d..366b6d23d9b67 100644
--- a/mypyc/test-data/run-strings.test
+++ b/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