Repository: python/mypy
Pull request: [mypyc] Pre-allocate space in list comprehensions
### Description

Closes https://github.com/mypyc/mypyc/issues/473

* Reconstruct `tuple_creation_from_generator_helper` to `preallocate_space_helper`
* Rewrite `new_tuple_with_length`. Now we need only to pass in length.
* Add new primitive for list initialization.
* Support faster list comprehension in some simple situations.

## Test Plan




Problem statement:
Pre-allocate space in list comprehensions
For list comprehensions like `[f(x) for x in alist]` we can predict the size of the new list (it's `len(alist)`). Maybe we should allocate enough space for all items at the beginning of the list comprehension. This would have these benefits:
* Faster comprehensions (no need to realloc).
* Less memory use (no extra unused space at the end of the list).

There might also be slightly less/more memory fragmentation, depending on how the memory allocator works.

~~As a bonus, we could replace the list append op with a specialized op that does fewer checks, since we know that there will be enough space.~~

Hints from issue discussion:
This is only possible when iterating over something that has a length. For other iterables we'd need to fall back to the current implementation.

Optionally, we could also require that the length of the iterated sequence doesn't change during iteration. This would allow us to use a specialized op instead of a list append op to add items. We'd need to guard against changes in the length of the sequence and raise an exception if the length is changed.
The implementation lives in `mypyc.irbuild.for_helpers.translate_list_comprehension`.
More hints:
* We'd only support for this iterables with type `tuple_rprimitive` and `list_rprimitive` (`node_type` can be used to look up the rtype of an expression). Also we can do this for `RTuple` (fixed-length tuples) if we box them first, but that's pretty marginal.
* Use `mypyc.primitives.list_ops.new_list_op` to construct the list.
* Use `builtin_len` to calculate the length of the sequence.
* You can use `new_list_op()` in `ll_builder.py` as a model of how to initialize the list efficiently.
* There may be some trickiness around tagged integers (shifted left by 1) and C integers (unshifted).

In high-level pseudocode the result would look something like this:
```
n = len(alist)
result = make uninitialized list of length n
i = 0
while i < n:
    if len(alist) ! = n: raise error
    result->items[i] = f(alist[i])  # directly set the item, don't use a C API function
    i += 1
```

Commit message: Fix IR test

Patch to review:
diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py
index 73c6ece3d708..538269d0e0b6 100644
--- a/mypyc/irbuild/for_helpers.py
+++ b/mypyc/irbuild/for_helpers.py
@@ -16,14 +16,15 @@
 )
 from mypyc.ir.rtypes import (
     RType, is_short_int_rprimitive, is_list_rprimitive, is_sequence_rprimitive,
-    RTuple, is_dict_rprimitive, short_int_rprimitive, int_rprimitive
+    is_tuple_rprimitive, is_dict_rprimitive,
+    RTuple, short_int_rprimitive, int_rprimitive
 )
 from mypyc.primitives.registry import CFunctionDescription
 from mypyc.primitives.dict_ops import (
     dict_next_key_op, dict_next_value_op, dict_next_item_op, dict_check_size_op,
     dict_key_iter_op, dict_value_iter_op, dict_item_iter_op
 )
-from mypyc.primitives.list_ops import list_append_op, list_get_item_unsafe_op
+from mypyc.primitives.list_ops import list_append_op, list_get_item_unsafe_op, new_list_set_item_op
 from mypyc.primitives.set_ops import set_add_op
 from mypyc.primitives.generic_ops import iter_op, next_op
 from mypyc.primitives.exc_ops import no_err_occurred_op
@@ -131,7 +132,41 @@ def for_loop_helper_with_index(builder: IRBuilder, index: Lvalue, expr: Expressi
     builder.activate_block(exit_block)
 
 
+def preallocate_space_helper(builder: IRBuilder,
+                             gen: GeneratorExpr,
+                             empty_op_llbuilder: Callable[[Value, int], Value],
+                             set_item_op: CFunctionDescription) -> Optional[Value]:
+    """Currently we only optimize for simplest generator expression.
+
+    "... for index in list/tuple"
+    """
+    if len(gen.sequences) == 1 and len(gen.indices) == 1 and len(gen.condlists[0]) == 0:
+        rtype = builder.node_type(gen.sequences[0])
+        if is_list_rprimitive(rtype) or is_tuple_rprimitive(rtype):
+
+            length = builder.builder.builtin_len(builder.accept(gen.sequences[0]),
+                                                 gen.line, use_pyssize_t=True)
+            target_op = empty_op_llbuilder(length, gen.line)
+
+            def set_item(item_index: Value) -> None:
+                e = builder.accept(gen.left_expr)
+                builder.call_c(set_item_op, [target_op, item_index, e], gen.line)
+
+            for_loop_helper_with_index(builder, gen.indices[0], gen.sequences[0],
+                                       set_item, gen.line)
+
+            return target_op
+    return None
+
+
 def translate_list_comprehension(builder: IRBuilder, gen: GeneratorExpr) -> Value:
+    # Try simplest list comprehension, otherwise fall back to general one
+    val = preallocate_space_helper(builder, gen,
+                                   empty_op_llbuilder=builder.builder.new_list_op_with_length,
+                                   set_item_op=new_list_set_item_op)
+    if val is not None:
+        return val
+
     list_ops = builder.new_list_op([], gen.line)
     loop_params = list(zip(gen.indices, gen.sequences, gen.condlists))
 
diff --git a/mypyc/irbuild/ll_builder.py b/mypyc/irbuild/ll_builder.py
index 0a0a9df050f3..2ed074d6c541 100644
--- a/mypyc/irbuild/ll_builder.py
+++ b/mypyc/irbuild/ll_builder.py
@@ -891,6 +891,9 @@ def make_dict(self, key_value_pairs: Sequence[DictEntry], line: int) -> Value:
 
         return result
 
+    def new_list_op_with_length(self, length: Value, line: int) -> Value:
+        return self.call_c(new_list_op, [length], line)
+
     def new_list_op(self, values: List[Value], line: int) -> Value:
         length = Integer(len(values), c_pyssize_t_rprimitive, line)
         result_list = self.call_c(new_list_op, [length], line)
@@ -1124,14 +1127,13 @@ 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
+    def new_tuple_with_length(self, length: Value, line: int) -> Value:
+        """Generate a new empty tuple with length
 
         Args:
-            val: a list or tuple
+            length: desired length, whose type should be c_pyssize_t_rprimitive
             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
diff --git a/mypyc/irbuild/specialize.py b/mypyc/irbuild/specialize.py
index cfb338e8864e..8d22ddb2fccf 100644
--- a/mypyc/irbuild/specialize.py
+++ b/mypyc/irbuild/specialize.py
@@ -22,14 +22,14 @@
 )
 from mypyc.ir.rtypes import (
     RType, RTuple, str_rprimitive, list_rprimitive, dict_rprimitive, set_rprimitive,
-    bool_rprimitive, is_dict_rprimitive, is_list_rprimitive, is_tuple_rprimitive
+    bool_rprimitive, is_dict_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, translate_set_comprehension,
-    comprehension_helper, for_loop_helper_with_index,
+    comprehension_helper, preallocate_space_helper
 )
 
 
@@ -150,8 +150,11 @@ 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 len(expr.args) == 1 and callee.fullname == "builtins.tuple":
+                val = preallocate_space_helper(
+                    builder, expr.args[0],
+                    empty_op_llbuilder=builder.builder.new_tuple_with_length,
+                    set_item_op=new_tuple_set_item_op)
                 if val is not None:
                     return val
 
@@ -162,28 +165,6 @@ def translate_safe_generator_call(
     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
diff --git a/mypyc/lib-rt/CPy.h b/mypyc/lib-rt/CPy.h
index 84b8bc81edec..0908af62bfb7 100644
--- a/mypyc/lib-rt/CPy.h
+++ b/mypyc/lib-rt/CPy.h
@@ -320,6 +320,7 @@ PyObject *CPyList_GetItem(PyObject *list, CPyTagged index);
 PyObject *CPyList_GetItemUnsafe(PyObject *list, CPyTagged index);
 PyObject *CPyList_GetItemShort(PyObject *list, CPyTagged index);
 bool CPyList_SetItem(PyObject *list, CPyTagged index, PyObject *value);
+bool CPyList_SetItemUnsafe(PyObject *list, CPyTagged index, PyObject *value);
 PyObject *CPyList_PopLast(PyObject *obj);
 PyObject *CPyList_Pop(PyObject *obj, CPyTagged index);
 CPyTagged CPyList_Count(PyObject *obj, PyObject *value);
diff --git a/mypyc/lib-rt/list_ops.c b/mypyc/lib-rt/list_ops.c
index 10706780d114..0030dcd403c7 100644
--- a/mypyc/lib-rt/list_ops.c
+++ b/mypyc/lib-rt/list_ops.c
@@ -79,7 +79,19 @@ bool CPyList_SetItem(PyObject *list, CPyTagged index, PyObject *value) {
         PyList_SET_ITEM(list, n, value);
         return true;
     } else {
-        PyErr_SetString(PyExc_IndexError, "list assignment index out of range");
+        PyErr_SetString(PyExc_OverflowError, "Python int too large to convert to C ssize_t");
+        return false;
+    }
+}
+
+// This function should only be used to fill in brand new lists.
+bool CPyList_SetItemUnsafe(PyObject *list, CPyTagged index, PyObject *value) {
+    if (CPyTagged_CheckShort(index)) {
+        Py_ssize_t n = CPyTagged_ShortAsSsize_t(index);
+        PyList_SET_ITEM(list, n, value);
+        return true;
+    } else {
+        PyErr_SetString(PyExc_OverflowError, "Python int too large to convert to C ssize_t");
         return false;
     }
 }
diff --git a/mypyc/primitives/list_ops.py b/mypyc/primitives/list_ops.py
index 3aae2099239e..5a541e453df2 100644
--- a/mypyc/primitives/list_ops.py
+++ b/mypyc/primitives/list_ops.py
@@ -65,6 +65,15 @@
     error_kind=ERR_FALSE,
     steals=[False, False, True])
 
+# PyList_SET_ITEM does no error checking,
+# and should only be used to fill in brand new tuples.
+new_list_set_item_op = custom_op(
+    arg_types=[list_rprimitive, int_rprimitive, object_rprimitive],
+    return_type=bit_rprimitive,
+    c_function_name='CPyList_SetItemUnsafe',
+    error_kind=ERR_FALSE,
+    steals=[False, False, True])
+
 # list.append(obj)
 list_append_op = method_op(
     name='append',
diff --git a/mypyc/test-data/irbuild-basic.test b/mypyc/test-data/irbuild-basic.test
index 9b7eaeb8e58b..40fe1f5aa809 100644
--- a/mypyc/test-data/irbuild-basic.test
+++ b/mypyc/test-data/irbuild-basic.test
@@ -2097,19 +2097,20 @@ def f(l):
     r6 :: tuple[int, int, int]
     r7, x, r8, y, r9, z :: int
     r10 :: short_int
-    r11 :: list
-    r12 :: short_int
-    r13 :: ptr
-    r14 :: native_int
-    r15 :: short_int
-    r16 :: bit
-    r17 :: object
-    r18 :: tuple[int, int, int]
-    r19, x_2, r20, y_2, r21, z_2, r22, r23 :: int
-    r24 :: object
-    r25 :: int32
-    r26 :: bit
-    r27 :: short_int
+    r11 :: ptr
+    r12 :: native_int
+    r13 :: list
+    r14 :: short_int
+    r15 :: ptr
+    r16 :: native_int
+    r17 :: short_int
+    r18 :: bit
+    r19 :: object
+    r20 :: tuple[int, int, int]
+    r21, x_2, r22, y_2, r23, z_2, r24, r25 :: int
+    r26 :: object
+    r27 :: bit
+    r28 :: short_int
 L0:
     r0 = 0
 L1:
@@ -2133,35 +2134,37 @@ L3:
     r0 = r10
     goto L1
 L4:
-    r11 = PyList_New(0)
-    r12 = 0
+    r11 = get_element_ptr l ob_size :: PyVarObject
+    r12 = load_mem r11 :: native_int*
+    keep_alive l
+    r13 = PyList_New(r12)
+    r14 = 0
 L5:
-    r13 = get_element_ptr l ob_size :: PyVarObject
-    r14 = load_mem r13 :: native_int*
+    r15 = get_element_ptr l ob_size :: PyVarObject
+    r16 = load_mem r15 :: native_int*
     keep_alive l
-    r15 = r14 << 1
-    r16 = r12 < r15 :: signed
-    if r16 goto L6 else goto L8 :: bool
+    r17 = r16 << 1
+    r18 = r14 < r17 :: signed
+    if r18 goto L6 else goto L8 :: bool
 L6:
-    r17 = CPyList_GetItemUnsafe(l, r12)
-    r18 = unbox(tuple[int, int, int], r17)
-    r19 = r18[0]
-    x_2 = r19
-    r20 = r18[1]
-    y_2 = r20
-    r21 = r18[2]
-    z_2 = r21
-    r22 = CPyTagged_Add(x_2, y_2)
-    r23 = CPyTagged_Add(r22, z_2)
-    r24 = box(int, r23)
-    r25 = PyList_Append(r11, r24)
-    r26 = r25 >= 0 :: signed
+    r19 = CPyList_GetItemUnsafe(l, r14)
+    r20 = unbox(tuple[int, int, int], r19)
+    r21 = r20[0]
+    x_2 = r21
+    r22 = r20[1]
+    y_2 = r22
+    r23 = r20[2]
+    z_2 = r23
+    r24 = CPyTagged_Add(x_2, y_2)
+    r25 = CPyTagged_Add(r24, z_2)
+    r26 = box(int, r25)
+    r27 = CPyList_SetItemUnsafe(r13, r14, r26)
 L7:
-    r27 = r12 + 2
-    r12 = r27
+    r28 = r14 + 2
+    r14 = r28
     goto L5
 L8:
-    return r11
+    return r13
 
 [case testProperty]
 class PropertyHolder:
diff --git a/mypyc/test-data/irbuild-lists.test b/mypyc/test-data/irbuild-lists.test
index 85d8448c07f1..9c5bb5d23bf1 100644
--- a/mypyc/test-data/irbuild-lists.test
+++ b/mypyc/test-data/irbuild-lists.test
@@ -255,3 +255,95 @@ L0:
     r1 = CPyList_Insert(x, 0, r0)
     r2 = r1 >= 0 :: signed
     return 1
+
+[case testListBuiltFromGenerator]
+from typing import List
+def f(source: List[int]) -> None:
+    a = list(x + 1 for x in source)
+    b = [x + 1 for x in source]
+[out]
+def f(source):
+    source :: list
+    r0 :: ptr
+    r1 :: native_int
+    r2 :: list
+    r3 :: short_int
+    r4 :: ptr
+    r5 :: native_int
+    r6 :: short_int
+    r7 :: bit
+    r8 :: object
+    r9, x, r10 :: int
+    r11 :: object
+    r12 :: bit
+    r13 :: short_int
+    r14 :: object
+    r15, a :: list
+    r16 :: ptr
+    r17 :: native_int
+    r18 :: list
+    r19 :: short_int
+    r20 :: ptr
+    r21 :: native_int
+    r22 :: short_int
+    r23 :: bit
+    r24 :: object
+    r25, x_2, r26 :: int
+    r27 :: object
+    r28 :: bit
+    r29 :: short_int
+    b :: list
+L0:
+    r0 = get_element_ptr source ob_size :: PyVarObject
+    r1 = load_mem r0 :: native_int*
+    keep_alive source
+    r2 = PyList_New(r1)
+    r3 = 0
+L1:
+    r4 = get_element_ptr source ob_size :: PyVarObject
+    r5 = load_mem r4 :: native_int*
+    keep_alive source
+    r6 = r5 << 1
+    r7 = r3 < r6 :: signed
+    if r7 goto L2 else goto L4 :: bool
+L2:
+    r8 = CPyList_GetItemUnsafe(source, r3)
+    r9 = unbox(int, r8)
+    x = r9
+    r10 = CPyTagged_Add(x, 2)
+    r11 = box(int, r10)
+    r12 = CPyList_SetItemUnsafe(r2, r3, r11)
+L3:
+    r13 = r3 + 2
+    r3 = r13
+    goto L1
+L4:
+    r14 = PyObject_GetIter(r2)
+    r15 = PySequence_List(r14)
+    a = r15
+    r16 = get_element_ptr source ob_size :: PyVarObject
+    r17 = load_mem r16 :: native_int*
+    keep_alive source
+    r18 = PyList_New(r17)
+    r19 = 0
+L5:
+    r20 = get_element_ptr source ob_size :: PyVarObject
+    r21 = load_mem r20 :: native_int*
+    keep_alive source
+    r22 = r21 << 1
+    r23 = r19 < r22 :: signed
+    if r23 goto L6 else goto L8 :: bool
+L6:
+    r24 = CPyList_GetItemUnsafe(source, r19)
+    r25 = unbox(int, r24)
+    x_2 = r25
+    r26 = CPyTagged_Add(x_2, 2)
+    r27 = box(int, r26)
+    r28 = CPyList_SetItemUnsafe(r18, r19, r27)
+L7:
+    r29 = r19 + 2
+    r19 = r29
+    goto L5
+L8:
+    b = r18
+    return 1
diff --git a/mypyc/test-data/run-lists.test b/mypyc/test-data/run-lists.test
index 8c567b80f596..9c36aafbd68e 100644
--- a/mypyc/test-data/run-lists.test
+++ b/mypyc/test-data/run-lists.test
@@ -338,3 +338,16 @@ assert list_in_mixed(0.0)
 assert not list_in_mixed([1])
 assert not list_in_mixed(object)
 assert list_in_mixed(type)
+
+[case testListBuiltFromGenerator]
+def test() -> None:
+    source = ["a", "b", "c"]
+    a = list(x + "f2" for x in source)
+    b = [x + "f2" for x in source]
+    c = [x + "f2" for x in (y + "yy" for y in source)]
+    assert a == ["af2", "bf2", "cf2"]
+    assert b == ["af2", "bf2", "cf2"]
+    assert c == ["ayyf2", "byyf2", "cyyf2"]
+    source2 = [1, 2, 3]
+    d = [x * 2 for x in source2]
+    assert d == [2, 4, 6]


Code review comments:
diff_hunk: @@ -65,6 +65,15 @@
     error_kind=ERR_FALSE,
     steals=[False, False, True])
 
+# PyList_SET_ITEM does no error checking,
+# and should only be used to fill in brand new tuples.
path: mypyc/primitives/list_ops.py
text: @user1:
tuples -> lists

diff_hunk: @@ -79,7 +79,19 @@ bool CPyList_SetItem(PyObject *list, CPyTagged index, PyObject *value) {
         PyList_SET_ITEM(list, n, value);
         return true;
     } else {
-        PyErr_SetString(PyExc_IndexError, "list assignment index out of range");
+        PyErr_SetString(PyExc_OverflowError, "Python int too large to convert to C ssize_t");
path: mypyc/lib-rt/list_ops.c
text: @user1:
Brainstorming: "int too large for list index" could be a slightly better message? We generally want to avoid C concepts in error messages.

@author:
I modified several error messages and now they are much clearer.

diff_hunk: @@ -131,7 +132,41 @@ def for_loop_helper_with_index(builder: IRBuilder, index: Lvalue, expr: Expressi
     builder.activate_block(exit_block)
 
 
+def preallocate_space_helper(builder: IRBuilder,
path: mypyc/irbuild/for_helpers.py
text: @user1:
Let's try to make the name more descriptive. Maybe `sequence_from_generator_preallocate_helper`?