bpo-32856: Optimize the assignment idiom in comprehensions. (GH-16814)

Now `for y in [expr]` in comprehensions is as fast as a simple
assignment `y = expr`.

8 files changed, 145 insertions, 18 deletions

diff --git a/Doc/whatsnew/3.9.rst b/Doc/whatsnew/3.9.rst
index c040410..ec17984 100644
--- a/Doc/whatsnew/3.9.rst
+++ b/Doc/whatsnew/3.9.rst
@@ -315,6 +315,17 @@ case), and one used ``__VENV_NAME__`` instead.
 Optimizations
 =============
 
+* Optimized the idiom for assignment a temporary variable in comprehensions.
+  Now ``for y in [expr]`` in comprehensions is as fast as a simple assignment
+  ``y = expr``.  For example:
+
+     sums = [s for s in [0] for x in data for s in [s + x]]
+
+  Unlike to the ``:=`` operator this idiom does not leak a variable to the
+  outer scope.
+
+  (Contributed by Serhiy Storchaka in :issue:`32856`.)
+
 
 Build and C API Changes
 =======================
diff --git a/Lib/test/test_dictcomps.py b/Lib/test/test_dictcomps.py
index 927e310..16aa651 100644
--- a/Lib/test/test_dictcomps.py
+++ b/Lib/test/test_dictcomps.py
@@ -111,5 +111,22 @@ class DictComprehensionTest(unittest.TestCase):
         self.assertEqual(actual, expected)
         self.assertEqual(actual_calls, expected_calls)
 
+    def test_assignment_idiom_in_comprehensions(self):
+        expected = {1: 1, 2: 4, 3: 9, 4: 16}
+        actual = {j: j*j for i in range(4) for j in [i+1]}
+        self.assertEqual(actual, expected)
+        expected = {3: 2, 5: 6, 7: 12, 9: 20}
+        actual = {j+k: j*k for i in range(4) for j in [i+1] for k in [j+1]}
+        self.assertEqual(actual, expected)
+        expected = {3: 2, 5: 6, 7: 12, 9: 20}
+        actual = {j+k: j*k for i in range(4)  for j, k in [(i+1, i+2)]}
+        self.assertEqual(actual, expected)
+
+    def test_star_expression(self):
+        expected = {0: 0, 1: 1, 2: 4, 3: 9}
+        self.assertEqual({i: i*i for i in [*range(4)]}, expected)
+        self.assertEqual({i: i*i for i in (*range(4),)}, expected)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Lib/test/test_genexps.py b/Lib/test/test_genexps.py
index fd712bb..86e4e19 100644
--- a/Lib/test/test_genexps.py
+++ b/Lib/test/test_genexps.py
@@ -15,6 +15,22 @@ Test nesting with the inner expression dependent on the outer
     >>> list((i,j) for i in range(4) for j in range(i) )
     [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)]
 
+Test the idiom for temporary variable assignment in comprehensions.
+
+    >>> list((j*j for i in range(4) for j in [i+1]))
+    [1, 4, 9, 16]
+    >>> list((j*k for i in range(4) for j in [i+1] for k in [j+1]))
+    [2, 6, 12, 20]
+    >>> list((j*k for i in range(4) for j, k in [(i+1, i+2)]))
+    [2, 6, 12, 20]
+
+Not assignment
+
+    >>> list((i*i for i in [*range(4)]))
+    [0, 1, 4, 9]
+    >>> list((i*i for i in (*range(4),)))
+    [0, 1, 4, 9]
+
 Make sure the induction variable is not exposed
 
     >>> i = 20
diff --git a/Lib/test/test_listcomps.py b/Lib/test/test_listcomps.py
index ddb169f..62b3319 100644
--- a/Lib/test/test_listcomps.py
+++ b/Lib/test/test_listcomps.py
@@ -16,6 +16,22 @@ Test nesting with the inner expression dependent on the outer
     >>> [(i,j) for i in range(4) for j in range(i)]
     [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)]
 
+Test the idiom for temporary variable assignment in comprehensions.
+
+    >>> [j*j for i in range(4) for j in [i+1]]
+    [1, 4, 9, 16]
+    >>> [j*k for i in range(4) for j in [i+1] for k in [j+1]]
+    [2, 6, 12, 20]
+    >>> [j*k for i in range(4) for j, k in [(i+1, i+2)]]
+    [2, 6, 12, 20]
+
+Not assignment
+
+    >>> [i*i for i in [*range(4)]]
+    [0, 1, 4, 9]
+    >>> [i*i for i in (*range(4),)]
+    [0, 1, 4, 9]
+
 Make sure the induction variable is not exposed
 
     >>> i = 20
diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py
index 567e6a1..7913e91 100644
--- a/Lib/test/test_peepholer.py
+++ b/Lib/test/test_peepholer.py
@@ -495,6 +495,20 @@ class TestTranforms(BytecodeTestCase):
             return 6
         self.check_lnotab(f)
 
+    def test_assignment_idiom_in_comprehensions(self):
+        def listcomp():
+            return [y for x in a for y in [f(x)]]
+        self.assertEqual(count_instr_recursively(listcomp, 'FOR_ITER'), 1)
+        def setcomp():
+            return {y for x in a for y in [f(x)]}
+        self.assertEqual(count_instr_recursively(setcomp, 'FOR_ITER'), 1)
+        def dictcomp():
+            return {y: y for x in a for y in [f(x)]}
+        self.assertEqual(count_instr_recursively(dictcomp, 'FOR_ITER'), 1)
+        def genexpr():
+            return (y for x in a for y in [f(x)])
+        self.assertEqual(count_instr_recursively(genexpr, 'FOR_ITER'), 1)
+
 
 class TestBuglets(unittest.TestCase):
 
diff --git a/Lib/test/test_setcomps.py b/Lib/test/test_setcomps.py
index fb7cde0..ecc4fff 100644
--- a/Lib/test/test_setcomps.py
+++ b/Lib/test/test_setcomps.py
@@ -21,6 +21,22 @@ Test nesting with the inner expression dependent on the outer
     >>> list(sorted({(i,j) for i in range(4) for j in range(i)}))
     [(1, 0), (2, 0), (2, 1), (3, 0), (3, 1), (3, 2)]
 
+Test the idiom for temporary variable assignment in comprehensions.
+
+    >>> sorted({j*j for i in range(4) for j in [i+1]})
+    [1, 4, 9, 16]
+    >>> sorted({j*k for i in range(4) for j in [i+1] for k in [j+1]})
+    [2, 6, 12, 20]
+    >>> sorted({j*k for i in range(4) for j, k in [(i+1, i+2)]})
+    [2, 6, 12, 20]
+
+Not assignment
+
+    >>> sorted({i*i for i in [*range(4)]})
+    [0, 1, 4, 9]
+    >>> sorted({i*i for i in (*range(4),)})
+    [0, 1, 4, 9]
+
 Make sure the induction variable is not exposed
 
     >>> i = 20
diff --git a/Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst b/Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst
new file mode 100644
index 0000000..c1cd68f
--- /dev/null
+++ b/Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst
@@ -0,0 +1,3 @@
+Optimized the idiom for assignment a temporary variable in comprehensions.
+Now ``for y in [expr]`` in comprehensions is as fast as a simple assignment
+``y = expr``.
diff --git a/Python/compile.c b/Python/compile.c
index 04b8fe4..bf8c810 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -212,11 +212,13 @@ static int compiler_set_qualname(struct compiler *);
 static int compiler_sync_comprehension_generator(
                                       struct compiler *c,
                                       asdl_seq *generators, int gen_index,
+                                      int depth,
                                       expr_ty elt, expr_ty val, int type);
 
 static int compiler_async_comprehension_generator(
                                       struct compiler *c,
                                       asdl_seq *generators, int gen_index,
+                                      int depth,
                                       expr_ty elt, expr_ty val, int type);
 
 static PyCodeObject *assemble(struct compiler *, int addNone);
@@ -4343,22 +4345,24 @@ ex_call:
 static int
 compiler_comprehension_generator(struct compiler *c,
                                  asdl_seq *generators, int gen_index,
+                                 int depth,
                                  expr_ty elt, expr_ty val, int type)
 {
     comprehension_ty gen;
     gen = (comprehension_ty)asdl_seq_GET(generators, gen_index);
     if (gen->is_async) {
         return compiler_async_comprehension_generator(
-            c, generators, gen_index, elt, val, type);
+            c, generators, gen_index, depth, elt, val, type);
     } else {
         return compiler_sync_comprehension_generator(
-            c, generators, gen_index, elt, val, type);
+            c, generators, gen_index, depth, elt, val, type);
     }
 }
 
 static int
 compiler_sync_comprehension_generator(struct compiler *c,
                                       asdl_seq *generators, int gen_index,
+                                      int depth,
                                       expr_ty elt, expr_ty val, int type)
 {
     /* generate code for the iterator, then each of the ifs,
@@ -4386,12 +4390,38 @@ compiler_sync_comprehension_generator(struct compiler *c,
     }
     else {
         /* Sub-iter - calculate on the fly */
-        VISIT(c, expr, gen->iter);
-        ADDOP(c, GET_ITER);
+        /* Fast path for the temporary variable assignment idiom:
+             for y in [f(x)]
+         */
+        asdl_seq *elts;
+        switch (gen->iter->kind) {
+            case List_kind:
+                elts = gen->iter->v.List.elts;
+                break;
+            case Tuple_kind:
+                elts = gen->iter->v.Tuple.elts;
+                break;
+            default:
+                elts = NULL;
+        }
+        if (asdl_seq_LEN(elts) == 1) {
+            expr_ty elt = asdl_seq_GET(elts, 0);
+            if (elt->kind != Starred_kind) {
+                VISIT(c, expr, elt);
+                start = NULL;
+            }
+        }
+        if (start) {
+            VISIT(c, expr, gen->iter);
+            ADDOP(c, GET_ITER);
+        }
+    }
+    if (start) {
+        depth++;
+        compiler_use_next_block(c, start);
+        ADDOP_JREL(c, FOR_ITER, anchor);
+        NEXT_BLOCK(c);
     }
-    compiler_use_next_block(c, start);
-    ADDOP_JREL(c, FOR_ITER, anchor);
-    NEXT_BLOCK(c);
     VISIT(c, expr, gen->target);
 
     /* XXX this needs to be cleaned up...a lot! */
@@ -4405,7 +4435,7 @@ compiler_sync_comprehension_generator(struct compiler *c,
 
     if (++gen_index < asdl_seq_LEN(generators))
         if (!compiler_comprehension_generator(c,
-                                              generators, gen_index,
+                                              generators, gen_index, depth,
                                               elt, val, type))
         return 0;
 
@@ -4420,18 +4450,18 @@ compiler_sync_comprehension_generator(struct compiler *c,
             break;
         case COMP_LISTCOMP:
             VISIT(c, expr, elt);
-            ADDOP_I(c, LIST_APPEND, gen_index + 1);
+            ADDOP_I(c, LIST_APPEND, depth + 1);
             break;
         case COMP_SETCOMP:
             VISIT(c, expr, elt);
-            ADDOP_I(c, SET_ADD, gen_index + 1);
+            ADDOP_I(c, SET_ADD, depth + 1);
             break;
         case COMP_DICTCOMP:
             /* With '{k: v}', k is evaluated before v, so we do
                the same. */
             VISIT(c, expr, elt);
             VISIT(c, expr, val);
-            ADDOP_I(c, MAP_ADD, gen_index + 1);
+            ADDOP_I(c, MAP_ADD, depth + 1);
             break;
         default:
             return 0;
@@ -4440,8 +4470,10 @@ compiler_sync_comprehension_generator(struct compiler *c,
         compiler_use_next_block(c, skip);
     }
     compiler_use_next_block(c, if_cleanup);
-    ADDOP_JABS(c, JUMP_ABSOLUTE, start);
-    compiler_use_next_block(c, anchor);
+    if (start) {
+        ADDOP_JABS(c, JUMP_ABSOLUTE, start);
+        compiler_use_next_block(c, anchor);
+    }
 
     return 1;
 }
@@ -4449,6 +4481,7 @@ compiler_sync_comprehension_generator(struct compiler *c,
 static int
 compiler_async_comprehension_generator(struct compiler *c,
                                       asdl_seq *generators, int gen_index,
+                                      int depth,
                                       expr_ty elt, expr_ty val, int type)
 {
     comprehension_ty gen;
@@ -4492,9 +4525,10 @@ compiler_async_comprehension_generator(struct compiler *c,
         NEXT_BLOCK(c);
     }
 
+    depth++;
     if (++gen_index < asdl_seq_LEN(generators))
         if (!compiler_comprehension_generator(c,
-                                              generators, gen_index,
+                                              generators, gen_index, depth,
                                               elt, val, type))
         return 0;
 
@@ -4509,18 +4543,18 @@ compiler_async_comprehension_generator(struct compiler *c,
             break;
         case COMP_LISTCOMP:
             VISIT(c, expr, elt);
-            ADDOP_I(c, LIST_APPEND, gen_index + 1);
+            ADDOP_I(c, LIST_APPEND, depth + 1);
             break;
         case COMP_SETCOMP:
             VISIT(c, expr, elt);
-            ADDOP_I(c, SET_ADD, gen_index + 1);
+            ADDOP_I(c, SET_ADD, depth + 1);
             break;
         case COMP_DICTCOMP:
             /* With '{k: v}', k is evaluated before v, so we do
                the same. */
             VISIT(c, expr, elt);
             VISIT(c, expr, val);
-            ADDOP_I(c, MAP_ADD, gen_index + 1);
+            ADDOP_I(c, MAP_ADD, depth + 1);
             break;
         default:
             return 0;
@@ -4583,7 +4617,7 @@ compiler_comprehension(struct compiler *c, expr_ty e, int type,
         ADDOP_I(c, op, 0);
     }
 
-    if (!compiler_comprehension_generator(c, generators, 0, elt,
+    if (!compiler_comprehension_generator(c, generators, 0, 0, elt,
                                           val, type))
         goto error_in_scope;