From 32d6b5812fe4919d3eb3013c5c490d232d0c5efe Mon Sep 17 00:00:00 2001
From: Jukka Lehtosalo <jukka.lehtosalo@iki.fi>
Date: Tue, 8 Jun 2021 21:41:58 +0100
Subject: [PATCH] [mypyc] Fix evaluation of iterable in list comprehension
 twice

This could result in a crash if the second evaluation results
in a shorter list, such as in this example (besides being incorrect
overall):

```
a = [s for s in f.readlines()]
```

`f.readlines()` was called twice, resulting in an empty list on
the second call. This caused the list object to have a NULL item,
which is invalid.

This should fix `mypy --install-types` in compiled mode.
---
 mypyc/irbuild/for_helpers.py  | 13 +++++++------
 mypyc/test-data/run-misc.test | 17 ++++++++++++++++-
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/mypyc/irbuild/for_helpers.py b/mypyc/irbuild/for_helpers.py
index 8097b453b117..2e44ec3afaf5 100644
--- a/mypyc/irbuild/for_helpers.py
+++ b/mypyc/irbuild/for_helpers.py
@@ -88,7 +88,10 @@ 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,
+def for_loop_helper_with_index(builder: IRBuilder,
+                               index: Lvalue,
+                               expr: Expression,
+                               expr_reg: Value,
                                body_insts: Callable[[Value], None], line: int) -> None:
     """Generate IR for a sequence iteration.
 
@@ -101,7 +104,6 @@ def for_loop_helper_with_index(builder: IRBuilder, index: Lvalue, expr: Expressi
         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)
 
@@ -163,16 +165,15 @@ def sequence_from_generator_preallocate_helper(
     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)
+            sequence = builder.accept(gen.sequences[0])
+            length = builder.builder.builtin_len(sequence, 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],
+            for_loop_helper_with_index(builder, gen.indices[0], gen.sequences[0], sequence,
                                        set_item, gen.line)
 
             return target_op
diff --git a/mypyc/test-data/run-misc.test b/mypyc/test-data/run-misc.test
index 46fa8d0e176f..8bb6c1cd394d 100644
--- a/mypyc/test-data/run-misc.test
+++ b/mypyc/test-data/run-misc.test
@@ -485,6 +485,8 @@ print(native.x)
 77
 
 [case testComprehensions]
+from typing import List
+
 # A list comprehension
 l = [str(x) + "     " + str(y) + "   " + str(x*y) for x in range(10)
      if x != 6 if x != 5 for y in range(x) if y*x != 8]
@@ -498,6 +500,17 @@ def pred(x: int) -> bool:
 # eventually and will raise an exception.
 l2 = [x for x in range(10) if x <= 6 if pred(x)]
 
+src = ['x']
+
+def f() -> List[str]:
+    global src
+    res = src
+    src = []
+    return res
+
+l3 = [s for s in f()]
+l4 = [s for s in f()]
+
 # A dictionary comprehension
 d = {k: k*k for k in range(10) if k != 5 if k != 6}
 
@@ -506,10 +519,12 @@ s = {str(x) + "     " + str(y) + "   " + str(x*y) for x in range(10)
      if x != 6 if x != 5 for y in range(x) if y*x != 8}
 
 [file driver.py]
-from native import l, l2, d, s
+from native import l, l2, l3, l4, d, s
 for a in l:
     print(a)
 print(tuple(l2))
+assert l3 == ['x']
+assert l4 == []
 for k in sorted(d):
     print(k, d[k])
 for a in sorted(s):