[mypyc] Optimize truth value testing for strings (#10269)

pranavrajpal · web-flow · commit 1e96f1df48a6 · 2021-04-06T17:02:57.000+01:00
Add a new primitive to check for empty strings. Related to mypyc/mypyc#768.
diff --git a/mypyc/irbuild/ll_builder.py b/mypyc/irbuild/ll_builder.py
@@ -63,7 +63,7 @@
 )
 from mypyc.primitives.int_ops import int_comparison_op_mapping
 from mypyc.primitives.exc_ops import err_occurred_op, keep_propagating_op
-from mypyc.primitives.str_ops import unicode_compare
+from mypyc.primitives.str_ops import unicode_compare, str_check_if_true
 from mypyc.primitives.set_ops import new_set_op
 from mypyc.rt_subtype import is_runtime_subtype
 from mypyc.subtype import is_subtype
@@ -958,6 +958,8 @@ def add_bool_branch(self, value: Value, true: BasicBlock, false: BasicBlock) ->
             zero = Integer(0, short_int_rprimitive)
             self.compare_tagged_condition(value, zero, '!=', true, false, value.line)
             return
+        elif is_same_type(value.type, str_rprimitive):
+            value = self.call_c(str_check_if_true, [value], value.line)
         elif is_same_type(value.type, list_rprimitive):
             length = self.builtin_len(value, value.line)
             zero = Integer(0)
diff --git a/mypyc/lib-rt/str_ops.c b/mypyc/lib-rt/str_ops.c
@@ -102,3 +102,8 @@ PyObject *CPyStr_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end) {
     }
     return CPyObject_GetSlice(obj, start, end);
 }
+/* Check if the given string is true (i.e. it's length isn't zero) */
+bool CPyStr_IsTrue(PyObject *obj) {
+    Py_ssize_t length = PyUnicode_GET_LENGTH(obj);
+    return length != 0;
+}
diff --git 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, bool_rprimitive
+    c_int_rprimitive, pointer_rprimitive, bool_rprimitive, bit_rprimitive
 )
 from mypyc.primitives.registry import (
     method_op, binary_op, function_op,
@@ -126,3 +126,11 @@
     return_type=str_rprimitive,
     c_function_name="CPyStr_Replace",
     error_kind=ERR_MAGIC)
+
+# check if a string is true (isn't an empty string)
+str_check_if_true = custom_op(
+    arg_types=[str_rprimitive],
+    return_type=bit_rprimitive,
+    c_function_name="CPyStr_IsTrue",
+    error_kind=ERR_NEVER,
+)
diff --git a/mypyc/test-data/irbuild-basic.test b/mypyc/test-data/irbuild-basic.test
@@ -196,24 +196,20 @@ def f(x: object, y: object) -> str:
 def f(x, y):
     x, y :: object
     r0 :: str
-    r1 :: int32
-    r2 :: bit
-    r3 :: bool
-    r4, r5 :: str
+    r1 :: bit
+    r2, r3 :: str
 L0:
     r0 = PyObject_Str(x)
-    r1 = PyObject_IsTrue(r0)
-    r2 = r1 >= 0 :: signed
-    r3 = truncate r1: int32 to builtins.bool
-    if r3 goto L1 else goto L2 :: bool
+    r1 = CPyStr_IsTrue(r0)
+    if r1 goto L1 else goto L2 :: bool
 L1:
-    r4 = r0
+    r2 = r0
     goto L3
 L2:
-    r5 = PyObject_Str(y)
-    r4 = r5
+    r3 = PyObject_Str(y)
+    r2 = r3
 L3:
-    return r4
+    return r2
 
 [case testOr]
 def f(x: int, y: int) -> int:
@@ -276,24 +272,20 @@ def f(x: object, y: object) -> str:
 def f(x, y):
     x, y :: object
     r0 :: str
-    r1 :: int32
-    r2 :: bit
-    r3 :: bool
-    r4, r5 :: str
+    r1 :: bit
+    r2, r3 :: str
 L0:
     r0 = PyObject_Str(x)
-    r1 = PyObject_IsTrue(r0)
-    r2 = r1 >= 0 :: signed
-    r3 = truncate r1: int32 to builtins.bool
-    if r3 goto L2 else goto L1 :: bool
+    r1 = CPyStr_IsTrue(r0)
+    if r1 goto L2 else goto L1 :: bool
 L1:
-    r4 = r0
+    r2 = r0
     goto L3
 L2:
-    r5 = PyObject_Str(y)
-    r4 = r5
+    r3 = PyObject_Str(y)
+    r2 = r3
 L3:
-    return r4
+    return r2
 
 [case testSimpleNot]
 def f(x: int, y: int) -> int:
diff --git a/mypyc/test-data/irbuild-statements.test b/mypyc/test-data/irbuild-statements.test
@@ -720,28 +720,24 @@ def complex_msg(x, s):
     r0 :: object
     r1 :: bit
     r2 :: str
-    r3 :: int32
-    r4 :: bit
-    r5 :: bool
-    r6 :: object
-    r7 :: str
-    r8, r9 :: object
+    r3 :: bit
+    r4 :: object
+    r5 :: str
+    r6, r7 :: object
 L0:
     r0 = load_address _Py_NoneStruct
     r1 = x != r0
     if r1 goto L1 else goto L2 :: bool
 L1:
     r2 = cast(str, x)
-    r3 = PyObject_IsTrue(r2)
-    r4 = r3 >= 0 :: signed
-    r5 = truncate r3: int32 to builtins.bool
-    if r5 goto L3 else goto L2 :: bool
+    r3 = CPyStr_IsTrue(r2)
+    if r3 goto L3 else goto L2 :: bool
 L2:
-    r6 = builtins :: module
-    r7 = 'AssertionError'
-    r8 = CPyObject_GetAttr(r6, r7)
-    r9 = PyObject_CallFunctionObjArgs(r8, s, 0)
-    CPy_Raise(r9)
+    r4 = builtins :: module
+    r5 = 'AssertionError'
+    r6 = CPyObject_GetAttr(r4, r5)
+    r7 = PyObject_CallFunctionObjArgs(r6, s, 0)
+    CPy_Raise(r7)
     unreachable
 L3:
     return 1
diff --git a/mypyc/test-data/irbuild-str.test b/mypyc/test-data/irbuild-str.test
@@ -138,4 +138,24 @@ L4:
     return r6 
 L5: 
     unreachable 
-    
+    
+[case testStrToBool]
+def is_true(x: str) -> bool:
+    if x:
+        return True
+    else:
+        return False
+[out]
+def is_true(x):
+    x :: str
+    r0 :: bit
+L0:
+    r0 = CPyStr_IsTrue(x)
+    if r0 goto L1 else goto L2 :: bool
+L1:
+    return 1
+L2:
+    return 0
+L3:
+    unreachable
+
diff --git a/mypyc/test-data/run-strings.test b/mypyc/test-data/run-strings.test
@@ -159,3 +159,22 @@ def test_str_replace() -> None:
     assert a.replace("foo", "bar", 4) == "barbarbar"
     assert a.replace("aaa", "bar") == "foofoofoo"
     assert a.replace("ofo", "xyzw") == "foxyzwxyzwo"
+
+def is_true(x: str) -> bool:
+    if x:
+        return True
+    else:
+        return False
+
+def is_false(x: str) -> bool:
+    if not x:
+        return True
+    else:
+        return False
+
+def test_str_to_bool() -> None:
+    assert is_false('')
+    assert not is_true('')
+    for x in 'a', 'foo', 'bar', 'some string':
+        assert is_true(x)
+        assert not is_false(x)
