[mypyc] Support list of extra constants and str.split (python#9334)

This PR changes the CallC description so that a list of extra constants can be used to 
a call. With this added functionality, str.split is implemented in new style IR.

mypyc/irbuild/ll_builder.py

@@ -777,8 +777,8 @@ def call_c(self,
                 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
+        for item in desc.extra_int_constants:
+            val, typ = item
             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,

mypyc/primitives/generic_ops.py

@@ -32,7 +32,7 @@
                 return_type=object_rprimitive,
                 c_function_name='PyObject_RichCompare',
                 error_kind=ERR_MAGIC,
-                extra_int_constant=(opid, c_int_rprimitive),
+                extra_int_constants=[(opid, c_int_rprimitive)],
                 priority=0)
 
 for op, funcname in [('+', 'PyNumber_Add'),
@@ -193,7 +193,7 @@
     c_function_name='PyObject_CallFunctionObjArgs',
     error_kind=ERR_MAGIC,
     var_arg_type=object_rprimitive,
-    extra_int_constant=(0, pointer_rprimitive))
+    extra_int_constants=[(0, pointer_rprimitive)])
 
 # Call callable object with positional + keyword args: func(*args, **kwargs)
 # Arguments are (func, *args tuple, **kwargs dict).
@@ -211,7 +211,7 @@
     c_function_name='CPyObject_CallMethodObjArgs',
     error_kind=ERR_MAGIC,
     var_arg_type=object_rprimitive,
-    extra_int_constant=(0, pointer_rprimitive))
+    extra_int_constants=[(0, pointer_rprimitive)])
 
 # len(obj)
 generic_len_op = c_custom_op(

mypyc/primitives/registry.py

@@ -52,7 +52,7 @@
                               ('error_kind', int),
                               ('steals', StealsDescription),
                               ('ordering', Optional[List[int]]),
-                              ('extra_int_constant', Optional[Tuple[int, RType]]),
+                              ('extra_int_constants', List[Tuple[int, RType]]),
                               ('priority', int)])
 
 # A description for C load operations including LoadGlobal and LoadAddress
@@ -307,7 +307,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,
+                extra_int_constants: List[Tuple[int, RType]] = [],
                 steals: StealsDescription = False,
                 priority: int = 1) -> CFunctionDescription:
     """Define a c function call op that replaces a method call.
@@ -329,13 +329,13 @@ 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
+        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)
         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_constant,
+                                c_function_name, error_kind, steals, ordering, extra_int_constants,
                                 priority)
     ops.append(desc)
     return desc
@@ -349,7 +349,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,
+                  extra_int_constants: List[Tuple[int, RType]] = [],
                   steals: StealsDescription = False,
                   priority: int = 1) -> CFunctionDescription:
     """Define a c function call op that replaces a function call.
@@ -364,7 +364,7 @@ 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_constant,
+                                c_function_name, error_kind, steals, ordering, extra_int_constants,
                                 priority)
     ops.append(desc)
     return desc
@@ -378,7 +378,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,
+                extra_int_constants: List[Tuple[int, RType]] = [],
                 steals: StealsDescription = False,
                 priority: int = 1) -> CFunctionDescription:
     """Define a c function call op for a binary operation.
@@ -390,7 +390,7 @@ 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_constant,
+                                c_function_name, error_kind, steals, ordering, extra_int_constants,
                                 priority)
     ops.append(desc)
     return desc
@@ -403,15 +403,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,
+                extra_int_constants: List[Tuple[int, RType]] = [],
                 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,
-                                extra_int_constant, 0)
+                                extra_int_constants, 0)
 
 
 def c_unary_op(name: str,
@@ -421,7 +421,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,
+               extra_int_constants: List[Tuple[int, RType]] = [],
                steals: StealsDescription = False,
                priority: int = 1) -> CFunctionDescription:
     """Define a c function call op for an unary operation.
@@ -433,7 +433,7 @@ 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_constant,
+                                c_function_name, error_kind, steals, ordering, extra_int_constants,
                                 priority)
     ops.append(desc)
     return desc

mypyc/primitives/str_ops.py

@@ -1,13 +1,14 @@
 """Primitive str ops."""
 
-from typing import List, Callable
+from typing import List, Callable, Tuple
 
-from mypyc.ir.ops import ERR_MAGIC, EmitterInterface, EmitCallback
+from mypyc.ir.ops import ERR_MAGIC, EmitterInterface
 from mypyc.ir.rtypes import (
-    RType, object_rprimitive, str_rprimitive, bool_rprimitive, int_rprimitive, list_rprimitive
+    RType, object_rprimitive, str_rprimitive, bool_rprimitive, int_rprimitive, list_rprimitive,
+    c_int_rprimitive, pointer_rprimitive
 )
 from mypyc.primitives.registry import (
-    binary_op, simple_emit, method_op, call_emit, c_method_op, c_binary_op, c_function_op,
+    binary_op, c_method_op, c_binary_op, c_function_op,
     load_address_op
 )
 
@@ -53,17 +54,19 @@
 
 # str.split(...)
 str_split_types = [str_rprimitive, str_rprimitive, int_rprimitive]  # type: List[RType]
-str_split_emits = [simple_emit('{dest} = PyUnicode_Split({args[0]}, NULL, -1);'),
-                   simple_emit('{dest} = PyUnicode_Split({args[0]}, {args[1]}, -1);'),
-                   call_emit('CPyStr_Split')] \
-                   # type: List[EmitCallback]
+str_split_functions = ['PyUnicode_Split', 'PyUnicode_Split', 'CPyStr_Split']
+str_split_constants = [[(0, pointer_rprimitive), (-1, c_int_rprimitive)],
+                       [(-1, c_int_rprimitive)],
+                       []] \
+                       # type: List[List[Tuple[int, RType]]]
 for i in range(len(str_split_types)):
-    method_op(
+    c_method_op(
         name='split',
         arg_types=str_split_types[0:i+1],
-        result_type=list_rprimitive,
-        error_kind=ERR_MAGIC,
-        emit=str_split_emits[i])
+        return_type=list_rprimitive,
+        c_function_name=str_split_functions[i],
+        extra_int_constants=str_split_constants[i],
+        error_kind=ERR_MAGIC)
 
 # str1 += str2
 #

mypyc/test-data/irbuild-str.test

@@ -0,0 +1,57 @@
+[case testStrSplit]
+from typing import Optional, List
+
+def do_split(s: str, sep: Optional[str] = None, max_split: Optional[int] = None) -> List[str]:
+    if sep is not None:
+        if max_split is not None:
+            return s.split(sep, max_split)
+        else:
+            return s.split(sep)
+    return s.split()
+[out]
+def do_split(s, sep, max_split):
+    s :: str
+    sep :: union[str, None]
+    max_split :: union[int, None]
+    r0, r1, r2 :: object
+    r3, r4 :: bool
+    r5 :: object
+    r6, r7 :: bool
+    r8 :: str
+    r9 :: int
+    r10 :: list
+    r11 :: str
+    r12, r13 :: list
+L0:
+    if is_error(sep) goto L1 else goto L2
+L1:
+    r0 = box(None, 1)
+    sep = r0
+L2:
+    if is_error(max_split) goto L3 else goto L4
+L3:
+    r1 = box(None, 1)
+    max_split = r1
+L4:
+    r2 = box(None, 1)
+    r3 = sep == r2
+    r4 = !r3
+    if r4 goto L5 else goto L9 :: bool
+L5:
+    r5 = box(None, 1)
+    r6 = max_split == r5
+    r7 = !r6
+    if r7 goto L6 else goto L7 :: bool
+L6:
+    r8 = cast(str, sep)
+    r9 = unbox(int, max_split)
+    r10 = CPyStr_Split(s, r8, r9)
+    return r10
+L7:
+    r11 = cast(str, sep)
+    r12 = PyUnicode_Split(s, r11, -1)
+    return r12
+L8:
+L9:
+    r13 = PyUnicode_Split(s, 0, -1)
+    return r13

mypyc/test/test_irbuild.py

@@ -27,6 +27,7 @@
     'irbuild-generics.test',
     'irbuild-try.test',
     'irbuild-set.test',
+    'irbuild-str.test',
     'irbuild-strip-asserts.test',
     'irbuild-int.test',
 ]