[mypyc] Merge string equality ops (#9363)

TH3CHARLie · web-flow · commit 53f12531e5cd · 2020-08-28T00:05:52.000+08:00
This PR merges str `==` and `!=` by directly building them in irbuild.

The old primitive relies on `ERR_MAGIC`  to handle the exception, now we use an `err_occurred_op` with `keep_propagating_op` to represent the same semantics.

Actually, our first several commits to replace `PrimitiveOp` with `CallC` cause an incorrect primitive lookup logic that at the point the generic compare is merged, string primitives would just use the generic ones. This PR will also fix this since we can obsolete the old binary op registry.
diff --git a/mypyc/ir/ops.py b/mypyc/ir/ops.py
@@ -1166,6 +1166,7 @@ def __init__(self,
                  args: List[Value],
                  ret_type: RType,
                  steals: StealsDescription,
+                 is_borrowed: bool,
                  error_kind: int,
                  line: int,
                  var_arg_idx: int = -1) -> None:
@@ -1175,6 +1176,7 @@ def __init__(self,
         self.args = args
         self.type = ret_type
         self.steals = steals
+        self.is_borrowed = is_borrowed
         self.var_arg_idx = var_arg_idx  # the position of the first variable argument in args
 
     def to_str(self, env: Environment) -> str:
diff --git a/mypyc/irbuild/ll_builder.py b/mypyc/irbuild/ll_builder.py
@@ -29,7 +29,8 @@
     bool_rprimitive, list_rprimitive, str_rprimitive, is_none_rprimitive, object_rprimitive,
     c_pyssize_t_rprimitive, is_short_int_rprimitive, is_tagged, PyVarObject, short_int_rprimitive,
     is_list_rprimitive, is_tuple_rprimitive, is_dict_rprimitive, is_set_rprimitive, PySetObject,
-    none_rprimitive, RTuple, is_bool_rprimitive
+    none_rprimitive, RTuple, is_bool_rprimitive, is_str_rprimitive, c_int_rprimitive,
+    pointer_rprimitive
 )
 from mypyc.ir.func_ir import FuncDecl, FuncSignature
 from mypyc.ir.class_ir import ClassIR, all_concrete_classes
@@ -38,7 +39,7 @@
     STATIC_PREFIX
 )
 from mypyc.primitives.registry import (
-    binary_ops, method_ops, func_ops,
+    method_ops, func_ops,
     c_method_call_ops, CFunctionDescription, c_function_ops,
     c_binary_ops, c_unary_ops
 )
@@ -56,6 +57,8 @@
     none_object_op, fast_isinstance_op, bool_op, type_is_op
 )
 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.rt_subtype import is_runtime_subtype
 from mypyc.subtype import is_subtype
 from mypyc.sametype import is_same_type
@@ -567,15 +570,15 @@ def binary_op(self,
         if expr_op in ('is', 'is not'):
             return self.translate_is_op(lreg, rreg, expr_op, line)
 
+        if (is_str_rprimitive(lreg.type) and is_str_rprimitive(rreg.type)
+                and expr_op in ('==', '!=')):
+            return self.compare_strings(lreg, rreg, expr_op, line)
+
         if is_tagged(lreg.type) and is_tagged(rreg.type) and expr_op in int_comparison_op_mapping:
             return self.compare_tagged(lreg, rreg, expr_op, line)
 
         call_c_ops_candidates = c_binary_ops.get(expr_op, [])
         target = self.matching_call_c(call_c_ops_candidates, [lreg, rreg], line)
-        if target:
-            return target
-        ops = binary_ops.get(expr_op, [])
-        target = self.matching_primitive_op(ops, [lreg, rreg], line)
         assert target, 'Unsupported binary operation: %s' % expr_op
         return target
 
@@ -626,6 +629,32 @@ def compare_tagged(self, lhs: Value, rhs: Value, op: str, line: int) -> Value:
         self.goto_and_activate(out)
         return result
 
+    def compare_strings(self, lhs: Value, rhs: Value, op: str, line: int) -> Value:
+        """Compare two strings"""
+        compare_result = self.call_c(unicode_compare, [lhs, rhs], line)
+        error_constant = self.add(LoadInt(-1, line, c_int_rprimitive))
+        compare_error_check = self.add(ComparisonOp(compare_result,
+                                                    error_constant, ComparisonOp.EQ, line))
+        exception_check, propagate, final_compare = BasicBlock(), BasicBlock(), BasicBlock()
+        branch = Branch(compare_error_check, exception_check, final_compare, Branch.BOOL_EXPR)
+        branch.negated = False
+        self.add(branch)
+        self.activate_block(exception_check)
+        check_error_result = self.call_c(err_occurred_op, [], line)
+        null = self.add(LoadInt(0, line, pointer_rprimitive))
+        compare_error_check = self.add(ComparisonOp(check_error_result,
+                                                    null, ComparisonOp.NEQ, line))
+        branch = Branch(compare_error_check, propagate, final_compare, Branch.BOOL_EXPR)
+        branch.negated = False
+        self.add(branch)
+        self.activate_block(propagate)
+        self.call_c(keep_propagating_op, [], line)
+        self.goto(final_compare)
+        self.activate_block(final_compare)
+        op_type = ComparisonOp.EQ if op == '==' else ComparisonOp.NEQ
+        return self.add(ComparisonOp(compare_result,
+                                     self.add(LoadInt(0, line, c_int_rprimitive)), op_type, line))
+
     def compare_tuples(self,
                        lhs: Value,
                        rhs: Value,
@@ -840,7 +869,7 @@ def call_c(self,
             extra_int_constant = self.add(LoadInt(val, line, rtype=typ))
             coerced.append(extra_int_constant)
         target = self.add(CallC(desc.c_function_name, coerced, desc.return_type, desc.steals,
-                                desc.error_kind, line, var_arg_idx))
+                                desc.is_borrowed, desc.error_kind, line, var_arg_idx))
         if desc.truncated_type is None:
             result = target
         else:
diff --git a/mypyc/primitives/exc_ops.py b/mypyc/primitives/exc_ops.py
@@ -41,6 +41,12 @@
     c_function_name='CPy_NoErrOccured',
     error_kind=ERR_FALSE)
 
+err_occurred_op = c_custom_op(
+    arg_types=[],
+    return_type=object_rprimitive,
+    c_function_name='PyErr_Occurred',
+    error_kind=ERR_NEVER,
+    is_borrowed=True)
 
 # Keep propagating a raised exception by unconditionally giving an error value.
 # This doesn't actually raise an exception.
diff --git a/mypyc/primitives/registry.py b/mypyc/primitives/registry.py
@@ -51,6 +51,7 @@
                               ('c_function_name', str),
                               ('error_kind', int),
                               ('steals', StealsDescription),
+                              ('is_borrowed', bool),
                               ('ordering', Optional[List[int]]),
                               ('extra_int_constants', List[Tuple[int, RType]]),
                               ('priority', int)])
@@ -61,9 +62,6 @@
                                    ('type', RType),
                                    ('src', str)])  # name of the target to load
 
-# Primitive binary ops (key is operator such as '+')
-binary_ops = {}  # type: Dict[str, List[OpDescription]]
-
 # Primitive ops for built-in functions (key is function name such as 'builtins.len')
 func_ops = {}  # type: Dict[str, List[OpDescription]]
 
@@ -116,31 +114,6 @@ def call_emit(func: str) -> EmitCallback:
     return simple_emit('{dest} = %s({comma_args});' % func)
 
 
-def binary_op(op: str,
-              arg_types: List[RType],
-              result_type: RType,
-              error_kind: int,
-              emit: EmitCallback,
-              format_str: Optional[str] = None,
-              steals: StealsDescription = False,
-              is_borrowed: bool = False,
-              priority: int = 1) -> None:
-    """Define a PrimitiveOp for a binary operation.
-
-    Arguments are similar to func_op(), but exactly two argument types
-    are expected.
-
-    This will be automatically generated by matching against the AST.
-    """
-    assert len(arg_types) == 2
-    ops = binary_ops.setdefault(op, [])
-    if format_str is None:
-        format_str = '{dest} = {args[0]} %s {args[1]}' % op
-    desc = OpDescription(op, arg_types, result_type, False, error_kind, format_str, emit,
-                         steals, is_borrowed, priority)
-    ops.append(desc)
-
-
 def func_op(name: str,
             arg_types: List[RType],
             result_type: RType,
@@ -281,6 +254,7 @@ def c_method_op(name: str,
                 ordering: Optional[List[int]] = None,
                 extra_int_constants: List[Tuple[int, RType]] = [],
                 steals: StealsDescription = False,
+                is_borrowed: bool = False,
                 priority: int = 1) -> CFunctionDescription:
     """Define a c function call op that replaces a method call.
 
@@ -303,12 +277,13 @@ def c_method_op(name: str,
                   accepted by the python syntax(before reordering)
         extra_int_constants: optional extra integer constants as the last arguments to a C call
         steals: description of arguments that this steals (ref count wise)
+        is_borrowed: if True, returned value is borrowed (no need to decrease refcount)
         priority: if multiple ops match, the one with the highest priority is picked
     """
     ops = c_method_call_ops.setdefault(name, [])
     desc = CFunctionDescription(name, arg_types, return_type, var_arg_type, truncated_type,
-                                c_function_name, error_kind, steals, ordering, extra_int_constants,
-                                priority)
+                                c_function_name, error_kind, steals, is_borrowed, ordering,
+                                extra_int_constants, priority)
     ops.append(desc)
     return desc
 
@@ -323,6 +298,7 @@ def c_function_op(name: str,
                   ordering: Optional[List[int]] = None,
                   extra_int_constants: List[Tuple[int, RType]] = [],
                   steals: StealsDescription = False,
+                  is_borrowed: bool = False,
                   priority: int = 1) -> CFunctionDescription:
     """Define a c function call op that replaces a function call.
 
@@ -336,8 +312,8 @@ def c_function_op(name: str,
     """
     ops = c_function_ops.setdefault(name, [])
     desc = CFunctionDescription(name, arg_types, return_type, var_arg_type, truncated_type,
-                                c_function_name, error_kind, steals, ordering, extra_int_constants,
-                                priority)
+                                c_function_name, error_kind, steals, is_borrowed, ordering,
+                                extra_int_constants, priority)
     ops.append(desc)
     return desc
 
@@ -352,6 +328,7 @@ def c_binary_op(name: str,
                 ordering: Optional[List[int]] = None,
                 extra_int_constants: List[Tuple[int, RType]] = [],
                 steals: StealsDescription = False,
+                is_borrowed: bool = False,
                 priority: int = 1) -> CFunctionDescription:
     """Define a c function call op for a binary operation.
 
@@ -362,8 +339,8 @@ def c_binary_op(name: str,
     """
     ops = c_binary_ops.setdefault(name, [])
     desc = CFunctionDescription(name, arg_types, return_type, var_arg_type, truncated_type,
-                                c_function_name, error_kind, steals, ordering, extra_int_constants,
-                                priority)
+                                c_function_name, error_kind, steals, is_borrowed, ordering,
+                                extra_int_constants, priority)
     ops.append(desc)
     return desc
 
@@ -376,13 +353,14 @@ def c_custom_op(arg_types: List[RType],
                 truncated_type: Optional[RType] = None,
                 ordering: Optional[List[int]] = None,
                 extra_int_constants: List[Tuple[int, RType]] = [],
-                steals: StealsDescription = False) -> CFunctionDescription:
+                steals: StealsDescription = False,
+                is_borrowed: bool = False) -> CFunctionDescription:
     """Create a one-off CallC op that can't be automatically generated from the AST.
 
     Most arguments are similar to c_method_op().
     """
     return CFunctionDescription('<custom>', arg_types, return_type, var_arg_type, truncated_type,
-                                c_function_name, error_kind, steals, ordering,
+                                c_function_name, error_kind, steals, is_borrowed, ordering,
                                 extra_int_constants, 0)
 
 
@@ -395,6 +373,7 @@ def c_unary_op(name: str,
                ordering: Optional[List[int]] = None,
                extra_int_constants: List[Tuple[int, RType]] = [],
                steals: StealsDescription = False,
+               is_borrowed: bool = False,
                priority: int = 1) -> CFunctionDescription:
     """Define a c function call op for an unary operation.
 
@@ -405,8 +384,8 @@ def c_unary_op(name: str,
     """
     ops = c_unary_ops.setdefault(name, [])
     desc = CFunctionDescription(name, [arg_type], return_type, None, truncated_type,
-                                c_function_name, error_kind, steals, ordering, extra_int_constants,
-                                priority)
+                                c_function_name, error_kind, steals, is_borrowed, ordering,
+                                extra_int_constants, priority)
     ops.append(desc)
     return desc
 
diff --git a/mypyc/primitives/str_ops.py b/mypyc/primitives/str_ops.py
@@ -1,15 +1,15 @@
 """Primitive str ops."""
 
-from typing import List, Callable, Tuple
+from typing import List, Tuple
 
-from mypyc.ir.ops import ERR_MAGIC, EmitterInterface
+from mypyc.ir.ops import ERR_MAGIC, ERR_NEVER
 from mypyc.ir.rtypes import (
-    RType, object_rprimitive, str_rprimitive, bool_rprimitive, int_rprimitive, list_rprimitive,
+    RType, object_rprimitive, str_rprimitive, int_rprimitive, list_rprimitive,
     c_int_rprimitive, pointer_rprimitive
 )
 from mypyc.primitives.registry import (
-    binary_op, c_method_op, c_binary_op, c_function_op,
-    load_address_op
+    c_method_op, c_binary_op, c_function_op,
+    load_address_op, c_custom_op
 )
 
 
@@ -80,30 +80,8 @@
             steals=[True, False])
 
 
-def emit_str_compare(comparison: str) -> Callable[[EmitterInterface, List[str], str], None]:
-    def emit(emitter: EmitterInterface, args: List[str], dest: str) -> None:
-        temp = emitter.temp_name()
-        emitter.emit_declaration('int %s;' % temp)
-        emitter.emit_lines(
-            '%s = PyUnicode_Compare(%s, %s);' % (temp, args[0], args[1]),
-            'if (%s == -1 && PyErr_Occurred())' % temp,
-            '    %s = 2;' % dest,
-            'else',
-            '    %s = (%s %s);' % (dest, temp, comparison))
-
-    return emit
-
-
-# str1 == str2
-binary_op(op='==',
-          arg_types=[str_rprimitive, str_rprimitive],
-          result_type=bool_rprimitive,
-          error_kind=ERR_MAGIC,
-          emit=emit_str_compare('== 0'))
-
-# str1 != str2
-binary_op(op='!=',
-          arg_types=[str_rprimitive, str_rprimitive],
-          result_type=bool_rprimitive,
-          error_kind=ERR_MAGIC,
-          emit=emit_str_compare('!= 0'))
+unicode_compare = c_custom_op(
+    arg_types=[str_rprimitive, str_rprimitive],
+    return_type=c_int_rprimitive,
+    c_function_name='PyUnicode_Compare',
+    error_kind=ERR_NEVER)
diff --git a/mypyc/test-data/irbuild-str.test b/mypyc/test-data/irbuild-str.test
@@ -55,3 +55,50 @@ L8:
 L9:
     r13 = PyUnicode_Split(s, 0, -1)
     return r13
+
+[case testStrEquality]
+def eq(x: str, y: str) -> bool:
+    return x == y
+
+def neq(x: str, y: str) -> bool:
+    return x != y
+
+[out]
+def eq(x, y):
+    x, y :: str
+    r0 :: int32
+    r1 :: bool
+    r2 :: object
+    r3, r4, r5 :: bool
+L0:
+    r0 = PyUnicode_Compare(x, y)
+    r1 = r0 == -1
+    if r1 goto L1 else goto L3 :: bool
+L1:
+    r2 = PyErr_Occurred()
+    r3 = r2 != 0
+    if r3 goto L2 else goto L3 :: bool
+L2:
+    r4 = CPy_KeepPropagating()
+L3:
+    r5 = r0 == 0
+    return r5
+def neq(x, y):
+    x, y :: str
+    r0 :: int32
+    r1 :: bool
+    r2 :: object
+    r3, r4, r5 :: bool
+L0:
+    r0 = PyUnicode_Compare(x, y)
+    r1 = r0 == -1
+    if r1 goto L1 else goto L3 :: bool
+L1:
+    r2 = PyErr_Occurred()
+    r3 = r2 != 0
+    if r3 goto L2 else goto L3 :: bool
+L2:
+    r4 = CPy_KeepPropagating()
+L3:
+    r5 = r0 != 0
+    return r5
diff --git a/mypyc/test/test_emitfunc.py b/mypyc/test/test_emitfunc.py
