[mypyc] Allow extra integer constant as the last argument to a C call (…

…#9251)

Related to mypyc/mypyc#734 and mypyc/mypyc#753.

mypyc/irbuild/ll_builder.py

@@ -761,6 +761,11 @@ def call_c(self,
                 arg = args[i]
                 arg = self.coerce(arg, desc.var_arg_type, line)
                 coerced.append(arg)
+        # add extra integer constant if any
+        if desc.extra_int_constant is not None:
+            val, typ = desc.extra_int_constant
+            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))
         if desc.truncated_type is None:

mypyc/primitives/generic_ops.py

@@ -10,28 +10,30 @@
 """
 
 from mypyc.ir.ops import ERR_NEVER, ERR_MAGIC, ERR_FALSE
-from mypyc.ir.rtypes import object_rprimitive, int_rprimitive, bool_rprimitive
+from mypyc.ir.rtypes import object_rprimitive, int_rprimitive, bool_rprimitive, c_int_rprimitive
 from mypyc.primitives.registry import (
     binary_op, unary_op, func_op, method_op, custom_op, call_emit, simple_emit,
-    call_negative_bool_emit, call_negative_magic_emit, negative_int_emit
+    call_negative_bool_emit, call_negative_magic_emit, negative_int_emit,
+    c_binary_op
 )
 
 
 # Binary operations
 
-for op, opid in [('==', 'Py_EQ'),
-                 ('!=', 'Py_NE'),
-                 ('<', 'Py_LT'),
-                 ('<=', 'Py_LE'),
-                 ('>', 'Py_GT'),
-                 ('>=', 'Py_GE')]:
+for op, opid in [('==', 2),   # PY_EQ
+                 ('!=', 3),   # PY_NE
+                 ('<',  0),   # PY_LT
+                 ('<=', 1),   # PY_LE
+                 ('>',  4),   # PY_GT
+                 ('>=', 5)]:  # PY_GE
     # The result type is 'object' since that's what PyObject_RichCompare returns.
-    binary_op(op=op,
-              arg_types=[object_rprimitive, object_rprimitive],
-              result_type=object_rprimitive,
-              error_kind=ERR_MAGIC,
-              emit=simple_emit('{dest} = PyObject_RichCompare({args[0]}, {args[1]}, %s);' % opid),
-              priority=0)
+    c_binary_op(name=op,
+                arg_types=[object_rprimitive, object_rprimitive],
+                return_type=object_rprimitive,
+                c_function_name='PyObject_RichCompare',
+                error_kind=ERR_MAGIC,
+                extra_int_constant=(opid, c_int_rprimitive),
+                priority=0)
 
 for op, funcname in [('+', 'PyNumber_Add'),
                      ('-', 'PyNumber_Subtract'),

mypyc/primitives/registry.py

@@ -35,7 +35,7 @@
 optimized implementations of all ops.
 """
 
-from typing import Dict, List, Optional, NamedTuple
+from typing import Dict, List, Optional, NamedTuple, Tuple
 
 from mypyc.ir.ops import (
     OpDescription, EmitterInterface, EmitCallback, StealsDescription, short_name
@@ -52,6 +52,7 @@
                               ('error_kind', int),
                               ('steals', StealsDescription),
                               ('ordering', Optional[List[int]]),
+                              ('extra_int_constant', Optional[Tuple[int, RType]]),
                               ('priority', int)])
 
 # A description for C load operations including LoadGlobal and LoadAddress
@@ -354,6 +355,7 @@ def c_method_op(name: str,
                 var_arg_type: Optional[RType] = None,
                 truncated_type: Optional[RType] = None,
                 ordering: Optional[List[int]] = None,
+                extra_int_constant: Optional[Tuple[int, RType]] = None,
                 steals: StealsDescription = False,
                 priority: int = 1) -> CFunctionDescription:
     """Define a c function call op that replaces a method call.
@@ -375,12 +377,14 @@ def c_method_op(name: str,
                   should never be used together with var_arg_type.
                   all the other arguments(such as arg_types) are in the order
                   accepted by the python syntax(before reordering)
+        extra_int_constant: optional extra integer constant as the last argument to a C call
         steals: description of arguments that this steals (ref count wise)
         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, priority)
+                                c_function_name, error_kind, steals, ordering, extra_int_constant,
+                                priority)
     ops.append(desc)
     return desc
 
@@ -393,6 +397,7 @@ def c_function_op(name: str,
                   var_arg_type: Optional[RType] = None,
                   truncated_type: Optional[RType] = None,
                   ordering: Optional[List[int]] = None,
+                  extra_int_constant: Optional[Tuple[int, RType]] = None,
                   steals: StealsDescription = False,
                   priority: int = 1) -> CFunctionDescription:
     """Define a c function call op that replaces a function call.
@@ -407,7 +412,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, priority)
+                                c_function_name, error_kind, steals, ordering, extra_int_constant,
+                                priority)
     ops.append(desc)
     return desc
 
@@ -420,6 +426,7 @@ def c_binary_op(name: str,
                 var_arg_type: Optional[RType] = None,
                 truncated_type: Optional[RType] = None,
                 ordering: Optional[List[int]] = None,
+                extra_int_constant: Optional[Tuple[int, RType]] = None,
                 steals: StealsDescription = False,
                 priority: int = 1) -> CFunctionDescription:
     """Define a c function call op for a binary operation.
@@ -431,7 +438,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, priority)
+                                c_function_name, error_kind, steals, ordering, extra_int_constant,
+                                priority)
     ops.append(desc)
     return desc
 
@@ -443,13 +451,15 @@ def c_custom_op(arg_types: List[RType],
                 var_arg_type: Optional[RType] = None,
                 truncated_type: Optional[RType] = None,
                 ordering: Optional[List[int]] = None,
+                extra_int_constant: Optional[Tuple[int, RType]] = None,
                 steals: StealsDescription = 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, 0)
+                                c_function_name, error_kind, steals, ordering,
+                                extra_int_constant, 0)
 
 
 def c_unary_op(name: str,
@@ -459,6 +469,7 @@ def c_unary_op(name: str,
                error_kind: int,
                truncated_type: Optional[RType] = None,
                ordering: Optional[List[int]] = None,
+               extra_int_constant: Optional[Tuple[int, RType]] = None,
                steals: StealsDescription = False,
                priority: int = 1) -> CFunctionDescription:
     """Define a c function call op for an unary operation.
@@ -470,7 +481,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, priority)
+                                c_function_name, error_kind, steals, ordering, extra_int_constant,
+                                priority)
     ops.append(desc)
     return desc
 

mypyc/test-data/irbuild-classes.test

@@ -855,15 +855,15 @@ def f(a, b):
     r0 :: object
     r1 :: bool
 L0:
-    r0 = a == b
+    r0 = PyObject_RichCompare(a, b, 2)
     r1 = unbox(bool, r0)
     return r1
 def f2(a, b):
     a, b :: __main__.Base
     r0 :: object
     r1 :: bool
 L0:
-    r0 = a != b
+    r0 = PyObject_RichCompare(a, b, 3)
     r1 = unbox(bool, r0)
     return r1
 def fOpt(a, b):
@@ -908,15 +908,15 @@ def f(a, b):
     r0 :: object
     r1 :: bool
 L0:
-    r0 = a == b
+    r0 = PyObject_RichCompare(a, b, 2)
     r1 = unbox(bool, r0)
     return r1
 def f2(a, b):
     a, b :: __main__.Base
     r0 :: object
     r1 :: bool
 L0:
-    r0 = a != b
+    r0 = PyObject_RichCompare(a, b, 3)
     r1 = unbox(bool, r0)
     return r1
 def fOpt(a, b):

mypyc/test/test_emitfunc.py

@@ -339,7 +339,7 @@ def test_simple(self) -> None:
         fn = FuncIR(FuncDecl('myfunc', None, 'mod', FuncSignature([self.arg], int_rprimitive)),
                     [self.block], self.env)
         emitter = Emitter(EmitterContext(NameGenerator([['mod']])))
-        generate_native_function(fn, emitter, 'prog.py', 'prog', False)
+        generate_native_function(fn, emitter, 'prog.py', 'prog', optimize_int=False)
         result = emitter.fragments
         assert_string_arrays_equal(
             [
@@ -358,7 +358,7 @@ def test_register(self) -> None:
         fn = FuncIR(FuncDecl('myfunc', None, 'mod', FuncSignature([self.arg], list_rprimitive)),
                     [self.block], self.env)
         emitter = Emitter(EmitterContext(NameGenerator([['mod']])))
-        generate_native_function(fn, emitter, 'prog.py', 'prog', False)
+        generate_native_function(fn, emitter, 'prog.py', 'prog', optimize_int=False)
         result = emitter.fragments
         assert_string_arrays_equal(
             [