[mypyc] Faster creation of a tuple from list/tuple (#10217)

Resolves mypyc/mypyc#769.

Author: 97littleleaf11

mypyc/irbuild/for_helpers.py

@@ -86,6 +86,50 @@ def for_loop_helper(builder: IRBuilder, index: Lvalue, expr: Expression,
     builder.activate_block(exit_block)
 
 
+def for_loop_helper_with_index(builder: IRBuilder, index: Lvalue, expr: Expression,
+                               body_insts: Callable[[Value], None], line: int) -> None:
+    """Generate IR for a sequence iteration.
+
+    This function only works for sequence type. Compared to for_loop_helper,
+    it would feed iteration index to body_insts.
+
+    Args:
+        index: the loop index Lvalue
+        expr: the expression to iterate over
+        body_insts: a function that generates the body of the loop.
+                    It needs a index as parameter.
+    """
+    expr_reg = builder.accept(expr)
+    assert is_sequence_rprimitive(expr_reg.type)
+    target_type = builder.get_sequence_type(expr)
+
+    body_block = BasicBlock()
+    step_block = BasicBlock()
+    exit_block = BasicBlock()
+    condition_block = BasicBlock()
+
+    for_gen = ForSequence(builder, index, body_block, exit_block, line, False)
+    for_gen.init(expr_reg, target_type, reverse=False)
+
+    builder.push_loop_stack(step_block, exit_block)
+
+    builder.goto_and_activate(condition_block)
+    for_gen.gen_condition()
+
+    builder.activate_block(body_block)
+    for_gen.begin_body()
+    body_insts(builder.read(for_gen.index_target))
+
+    builder.goto_and_activate(step_block)
+    for_gen.gen_step()
+    builder.goto(condition_block)
+
+    for_gen.add_cleanup(exit_block)
+    builder.pop_loop_stack()
+
+    builder.activate_block(exit_block)
+
+
 def translate_list_comprehension(builder: IRBuilder, gen: GeneratorExpr) -> Value:
     list_ops = builder.new_list_op([], gen.line)
     loop_params = list(zip(gen.indices, gen.sequences, gen.condlists))

mypyc/irbuild/ll_builder.py

@@ -47,13 +47,16 @@
 from mypyc.primitives.list_ops import (
     list_extend_op, new_list_op
 )
-from mypyc.primitives.tuple_ops import list_tuple_op, new_tuple_op
+from mypyc.primitives.tuple_ops import (
+    list_tuple_op, new_tuple_op, new_tuple_with_length_op
+)
 from mypyc.primitives.dict_ops import (
     dict_update_in_display_op, dict_new_op, dict_build_op, dict_size_op
 )
 from mypyc.primitives.generic_ops import (
-    py_getattr_op, py_call_op, py_call_with_kwargs_op, py_method_call_op, generic_len_op,
-    py_vectorcall_op, py_vectorcall_method_op
+    py_getattr_op, py_call_op, py_call_with_kwargs_op, py_method_call_op,
+    py_vectorcall_op, py_vectorcall_method_op,
+    generic_len_op, generic_ssize_t_len_op
 )
 from mypyc.primitives.misc_ops import (
     none_object_op, fast_isinstance_op, bool_op
@@ -1079,35 +1082,56 @@ def int_op(self, type: RType, lhs: Value, rhs: Value, op: int, line: int) -> Val
     def comparison_op(self, lhs: Value, rhs: Value, op: int, line: int) -> Value:
         return self.add(ComparisonOp(lhs, rhs, op, line))
 
-    def builtin_len(self, val: Value, line: int) -> Value:
+    def builtin_len(self, val: Value, line: int,
+                    use_pyssize_t: bool = False) -> Value:
+        """Return short_int_rprimitive by default."""
         typ = val.type
         if is_list_rprimitive(typ) or is_tuple_rprimitive(typ):
             elem_address = self.add(GetElementPtr(val, PyVarObject, 'ob_size'))
             size_value = self.add(LoadMem(c_pyssize_t_rprimitive, elem_address))
             self.add(KeepAlive([val]))
+            if use_pyssize_t:
+                return size_value
             offset = Integer(1, c_pyssize_t_rprimitive, line)
             return self.int_op(short_int_rprimitive, size_value, offset,
                                IntOp.LEFT_SHIFT, line)
         elif is_dict_rprimitive(typ):
             size_value = self.call_c(dict_size_op, [val], line)
+            if use_pyssize_t:
+                return size_value
             offset = Integer(1, c_pyssize_t_rprimitive, line)
             return self.int_op(short_int_rprimitive, size_value, offset,
                                IntOp.LEFT_SHIFT, line)
         elif is_set_rprimitive(typ):
             elem_address = self.add(GetElementPtr(val, PySetObject, 'used'))
             size_value = self.add(LoadMem(c_pyssize_t_rprimitive, elem_address))
             self.add(KeepAlive([val]))
+            if use_pyssize_t:
+                return size_value
             offset = Integer(1, c_pyssize_t_rprimitive, line)
             return self.int_op(short_int_rprimitive, size_value, offset,
                                IntOp.LEFT_SHIFT, line)
         # generic case
         else:
-            return self.call_c(generic_len_op, [val], line)
+            if use_pyssize_t:
+                return self.call_c(generic_ssize_t_len_op, [val], line)
+            else:
+                return self.call_c(generic_len_op, [val], line)
 
     def new_tuple(self, items: List[Value], line: int) -> Value:
         size = Integer(len(items), c_pyssize_t_rprimitive)  # type: Value
         return self.call_c(new_tuple_op, [size] + items, line)
 
+    def new_tuple_with_length(self, val: Value, line: int) -> Value:
+        """Generate a new empty tuple with length from a list or tuple
+
+        Args:
+            val: a list or tuple
+            line: line number
+        """
+        length = self.builtin_len(val, line, use_pyssize_t=True)
+        return self.call_c(new_tuple_with_length_op, [length], line)
+
     # Internal helpers
 
     def decompose_union_helper(self,

mypyc/irbuild/specialize.py

@@ -22,11 +22,15 @@
 )
 from mypyc.ir.rtypes import (
     RType, RTuple, str_rprimitive, list_rprimitive, dict_rprimitive, set_rprimitive,
-    bool_rprimitive, is_dict_rprimitive
+    bool_rprimitive, is_dict_rprimitive, is_list_rprimitive, is_tuple_rprimitive
 )
 from mypyc.primitives.dict_ops import dict_keys_op, dict_values_op, dict_items_op
+from mypyc.primitives.tuple_ops import new_tuple_set_item_op
 from mypyc.irbuild.builder import IRBuilder
-from mypyc.irbuild.for_helpers import translate_list_comprehension, comprehension_helper
+from mypyc.irbuild.for_helpers import (
+    translate_list_comprehension, comprehension_helper,
+    for_loop_helper_with_index
+)
 
 
 # Specializers are attempted before compiling the arguments to the
@@ -135,13 +139,39 @@ def translate_safe_generator_call(
                     + [builder.accept(arg) for arg in expr.args[1:]]),
                 builder.node_type(expr), expr.line, expr.arg_kinds, expr.arg_names)
         else:
+            if callee.fullname == "builtins.tuple":
+                val = tuple_from_generator_helper(builder, expr.args[0])
+                if val is not None:
+                    return val
             return builder.call_refexpr_with_args(
                 expr, callee,
                 ([translate_list_comprehension(builder, expr.args[0])]
                     + [builder.accept(arg) for arg in expr.args[1:]]))
     return None
 
 
+def tuple_from_generator_helper(builder: IRBuilder,
+                                gen: GeneratorExpr) -> Optional[Value]:
+
+    if len(gen.sequences) == 1 and len(gen.condlists[0]) == 0:
+        # Currently we only optimize for simplest generator expression
+        rtype = builder.node_type(gen.sequences[0])
+        if is_list_rprimitive(rtype) or is_tuple_rprimitive(rtype):
+            tuple_ops = builder.builder.new_tuple_with_length(builder.accept(gen.sequences[0]),
+                                                              gen.line)
+            item, expr = gen.indices[0], gen.sequences[0]
+
+            def set_tuple_item(item_index: Value) -> None:
+                e = builder.accept(gen.left_expr)
+                builder.call_c(new_tuple_set_item_op, [tuple_ops, item_index, e], gen.line)
+
+            for_loop_helper_with_index(builder, item, expr,
+                                       set_tuple_item, gen.line)
+
+            return tuple_ops
+    return None
+
+
 @specialize_function('builtins.any')
 def translate_any_call(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Optional[Value]:
     if (len(expr.args) == 1

mypyc/lib-rt/CPy.h

@@ -398,6 +398,7 @@ bool CPySet_Remove(PyObject *set, PyObject *key);
 
 PyObject *CPySequenceTuple_GetItem(PyObject *tuple, CPyTagged index);
 PyObject *CPySequenceTuple_GetSlice(PyObject *obj, CPyTagged start, CPyTagged end);
+bool CPySequenceTuple_SetItemUnsafe(PyObject *tuple, CPyTagged index, PyObject *value);
 
 
 // Exception operations

mypyc/lib-rt/tuple_ops.c

@@ -45,3 +45,17 @@ PyObject *CPySequenceTuple_GetSlice(PyObject *obj, CPyTagged start, CPyTagged en
     }
     return CPyObject_GetSlice(obj, start, end);
 }
+
+// PyTuple_SET_ITEM does no error checking,
+// and should only be used to fill in brand new tuples.
+bool CPySequenceTuple_SetItemUnsafe(PyObject *tuple, CPyTagged index, PyObject *value)
+{
+    if (CPyTagged_CheckShort(index)) {
+        Py_ssize_t n = CPyTagged_ShortAsSsize_t(index);
+        PyTuple_SET_ITEM(tuple, n, value);
+        return true;
+    } else {
+        PyErr_SetString(PyExc_OverflowError, "Python int too large to convert to C ssize_t");
+        return false;
+    }
+}

mypyc/primitives/generic_ops.py

@@ -12,7 +12,7 @@
 from mypyc.ir.ops import ERR_NEVER, ERR_MAGIC
 from mypyc.ir.rtypes import (
     object_rprimitive, int_rprimitive, bool_rprimitive, c_int_rprimitive, pointer_rprimitive,
-    object_pointer_rprimitive, c_size_t_rprimitive
+    object_pointer_rprimitive, c_size_t_rprimitive, c_pyssize_t_rprimitive
 )
 from mypyc.primitives.registry import (
     binary_op, c_unary_op, method_op, function_op, custom_op, ERR_NEG_INT
@@ -239,7 +239,15 @@
     arg_types=[object_rprimitive],
     return_type=int_rprimitive,
     c_function_name='CPyObject_Size',
-    error_kind=ERR_NEVER)
+    error_kind=ERR_MAGIC)
+
+# len(obj)
+# same as generic_len_op, however return py_ssize_t
+generic_ssize_t_len_op = custom_op(
+    arg_types=[object_rprimitive],
+    return_type=c_pyssize_t_rprimitive,
+    c_function_name='PyObject_Size',
+    error_kind=ERR_NEG_INT)
 
 # iter(obj)
 iter_op = function_op(name='builtins.iter',

mypyc/primitives/tuple_ops.py

@@ -4,9 +4,10 @@
 objects, i.e. tuple_rprimitive (RPrimitive), not RTuple.
 """
 
-from mypyc.ir.ops import ERR_MAGIC
+from mypyc.ir.ops import ERR_MAGIC, ERR_FALSE
 from mypyc.ir.rtypes import (
-    tuple_rprimitive, int_rprimitive, list_rprimitive, object_rprimitive, c_pyssize_t_rprimitive
+    tuple_rprimitive, int_rprimitive, list_rprimitive, object_rprimitive,
+    c_pyssize_t_rprimitive, bit_rprimitive
 )
 from mypyc.primitives.registry import method_op, function_op, custom_op
 
@@ -27,6 +28,21 @@
     error_kind=ERR_MAGIC,
     var_arg_type=object_rprimitive)
 
+new_tuple_with_length_op = custom_op(
+    arg_types=[c_pyssize_t_rprimitive],
+    return_type=tuple_rprimitive,
+    c_function_name='PyTuple_New',
+    error_kind=ERR_MAGIC)
+
+# PyTuple_SET_ITEM does no error checking,
+# and should only be used to fill in brand new tuples.
+new_tuple_set_item_op = custom_op(
+    arg_types=[tuple_rprimitive, int_rprimitive, object_rprimitive],
+    return_type=bit_rprimitive,
+    c_function_name='CPySequenceTuple_SetItemUnsafe',
+    error_kind=ERR_FALSE,
+    steals=[False, False, True])
+
 # Construct tuple from a list.
 list_tuple_op = function_op(
     name='builtins.tuple',