From 638cf8c815354028a96ac013cfada47d74bb546d Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Fri, 16 Feb 2018 10:44:39 +0200
Subject: [PATCH 1/6] bpo-32856: Optimiz the idiom for assignment in
 comprehensions.

Now `for y in [expr]` in comprehensions is so fast as a simple
assignment `y = expr`.
---
 .../2018-02-16-10-44-24.bpo-32856.UjR8SD.rst  |  3 ++
 Python/compile.c                              | 41 +++++++++++++++----
 2 files changed, 37 insertions(+), 7 deletions(-)
 create mode 100644 Misc/NEWS.d/next/Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst

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 00000000000000..4233a919807a28
--- /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 so fast as a simple assignment
+``y = expr``.
diff --git a/Python/compile.c b/Python/compile.c
index 0dc35662749b9c..1d42df62827bb1 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -3825,12 +3825,37 @@ 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) {
+        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! */
@@ -3879,8 +3904,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;
 }

From 59fb61772d184b2eaf8f7b41ee0f0a96e4981d7d Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Fri, 16 Feb 2018 19:16:04 +0200
Subject: [PATCH 2/6] Fix a bug and add tests.

---
 Lib/test/test_dictcomps.py | 11 +++++++++++
 Lib/test/test_genexps.py   |  9 +++++++++
 Lib/test/test_listcomps.py |  9 +++++++++
 Lib/test/test_peepholer.py | 27 +++++++++++++++++++++++++++
 Lib/test/test_setcomps.py  |  9 +++++++++
 Python/compile.c           | 29 ++++++++++++++++++-----------
 6 files changed, 83 insertions(+), 11 deletions(-)

diff --git a/Lib/test/test_dictcomps.py b/Lib/test/test_dictcomps.py
index 0873071883439e..bac5d3e23e1f0f 100644
--- a/Lib/test/test_dictcomps.py
+++ b/Lib/test/test_dictcomps.py
@@ -81,6 +81,17 @@ def test_illegal_assignment(self):
             compile("{x: y for y, x in ((1, 2), (3, 4))} += 5", "<test>",
                     "exec")
 
+    def test_assignment_idiom_in_comprehesions(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)
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Lib/test/test_genexps.py b/Lib/test/test_genexps.py
index fb531d6d472b6f..5305204900f361 100644
--- a/Lib/test/test_genexps.py
+++ b/Lib/test/test_genexps.py
@@ -15,6 +15,15 @@
     >>> 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]
+
 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 ddb169fe58957c..da441da7d28b07 100644
--- a/Lib/test/test_listcomps.py
+++ b/Lib/test/test_listcomps.py
@@ -16,6 +16,15 @@
     >>> [(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]
+
 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 0cc1e92907b52b..5f365f71a9e064 100644
--- a/Lib/test/test_peepholer.py
+++ b/Lib/test/test_peepholer.py
@@ -3,6 +3,19 @@
 
 from test.bytecode_helper import BytecodeTestCase
 
+def count_instr_recursively(f, opname):
+    count = 0
+    for instr in dis.get_instructions(f):
+        if instr.opname == opname:
+            count += 1
+    if hasattr(f, '__code__'):
+        f = f.__code__
+    for c in f.co_consts:
+        if hasattr(c, 'co_code'):
+            count += count_instr_recursively(c, opname)
+    return count
+
+
 class TestTranforms(BytecodeTestCase):
 
     def test_unot(self):
@@ -311,6 +324,20 @@ def test_constant_folding(self):
                 self.assertFalse(instr.opname.startswith('BINARY_'))
                 self.assertFalse(instr.opname.startswith('BUILD_'))
 
+    def test_assignment_idiom_in_comprehesions(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 fb7cde03d78236..ad954b04259b86 100644
--- a/Lib/test/test_setcomps.py
+++ b/Lib/test/test_setcomps.py
@@ -21,6 +21,15 @@
     >>> 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.
+
+    >>> list(sorted({j*j for i in range(4) for j in [i+1]}))
+    [1, 4, 9, 16]
+    >>> list(sorted({j*k for i in range(4) for j in [i+1] for k in [j+1]}))
+    [2, 6, 12, 20]
+    >>> list(sorted({j*k for i in range(4) for j, k in [(i+1, i+2)]}))
+    [2, 6, 12, 20]
+
 Make sure the induction variable is not exposed
 
     >>> i = 20
diff --git a/Python/compile.c b/Python/compile.c
index 1d42df62827bb1..24d5bfd4bfdea4 100644
--- a/Python/compile.c
+++ b/Python/compile.c
@@ -205,11 +205,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);
@@ -3782,22 +3784,24 @@ compiler_call_helper(struct compiler *c,
 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,
@@ -3852,6 +3856,7 @@ compiler_sync_comprehension_generator(struct compiler *c,
         }
     }
     if (start) {
+        depth++;
         compiler_use_next_block(c, start);
         ADDOP_JREL(c, FOR_ITER, anchor);
         NEXT_BLOCK(c);
@@ -3869,7 +3874,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;
 
@@ -3884,18 +3889,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 'd[k] = v', v is evaluated before k, so we do
                the same. */
             VISIT(c, expr, val);
             VISIT(c, expr, elt);
-            ADDOP_I(c, MAP_ADD, gen_index + 1);
+            ADDOP_I(c, MAP_ADD, depth + 1);
             break;
         default:
             return 0;
@@ -3915,6 +3920,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)
 {
     _Py_IDENTIFIER(StopAsyncIteration);
@@ -3998,9 +4004,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;
 
@@ -4015,18 +4022,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 'd[k] = v', v is evaluated before k, so we do
                the same. */
             VISIT(c, expr, val);
             VISIT(c, expr, elt);
-            ADDOP_I(c, MAP_ADD, gen_index + 1);
+            ADDOP_I(c, MAP_ADD, depth + 1);
             break;
         default:
             return 0;
@@ -4094,7 +4101,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;
 

From ae4af4865a8c032ef064fb50ebaeeaa84f298651 Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Fri, 23 Feb 2018 12:05:44 +0200
Subject: [PATCH 3/6] Fix typos.

---
 Lib/test/test_dictcomps.py | 2 +-
 Lib/test/test_peepholer.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/Lib/test/test_dictcomps.py b/Lib/test/test_dictcomps.py
index bac5d3e23e1f0f..23812006dde246 100644
--- a/Lib/test/test_dictcomps.py
+++ b/Lib/test/test_dictcomps.py
@@ -81,7 +81,7 @@ def test_illegal_assignment(self):
             compile("{x: y for y, x in ((1, 2), (3, 4))} += 5", "<test>",
                     "exec")
 
-    def test_assignment_idiom_in_comprehesions(self):
+    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)
diff --git a/Lib/test/test_peepholer.py b/Lib/test/test_peepholer.py
index 5f365f71a9e064..6484310888f1da 100644
--- a/Lib/test/test_peepholer.py
+++ b/Lib/test/test_peepholer.py
@@ -324,7 +324,7 @@ def test_constant_folding(self):
                 self.assertFalse(instr.opname.startswith('BINARY_'))
                 self.assertFalse(instr.opname.startswith('BUILD_'))
 
-    def test_assignment_idiom_in_comprehesions(self):
+    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)

From 8989629cf786b4e94dfd9742941bb38e0da7bfdd Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Wed, 16 Oct 2019 00:57:02 +0300
Subject: [PATCH 4/6] Add a What's New entry.

---
 Doc/whatsnew/3.9.rst | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/Doc/whatsnew/3.9.rst b/Doc/whatsnew/3.9.rst
index 6b4abc0567c30e..873532eb75a19c 100644
--- a/Doc/whatsnew/3.9.rst
+++ b/Doc/whatsnew/3.9.rst
@@ -158,9 +158,21 @@ now raises :exc:`ImportError` instead of :exc:`ValueError` for invalid relative
 import attempts.
 (Contributed by Ngalim Siregar in :issue:`37444`.)
 
+
 Optimizations
 =============
 
+* Optimized the idiom for assignment a temporary variable in comprehensions.
+  Now ``for y in [expr]`` in comprehensions is so 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
 =======================

From 3a53d9849d2b47ce2e91efc26449c165097b9300 Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Fri, 10 Jan 2020 10:41:02 +0200
Subject: [PATCH 5/6] Fix wording.

---
 Doc/whatsnew/3.9.rst                                            | 2 +-
 .../Core and Builtins/2018-02-16-10-44-24.bpo-32856.UjR8SD.rst  | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/Doc/whatsnew/3.9.rst b/Doc/whatsnew/3.9.rst
index cdea2575d91ab5..c1a8836633873e 100644
--- a/Doc/whatsnew/3.9.rst
+++ b/Doc/whatsnew/3.9.rst
@@ -238,7 +238,7 @@ Optimizations
 =============
 
 * Optimized the idiom for assignment a temporary variable in comprehensions.
-  Now ``for y in [expr]`` in comprehensions is so fast as a simple assignment
+  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]]
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
index 4233a919807a28..c1cd68f672712a 100644
--- 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	
@@ -1,3 +1,3 @@
 Optimized the idiom for assignment a temporary variable in comprehensions.
-Now ``for y in [expr]`` in comprehensions is so fast as a simple assignment
+Now ``for y in [expr]`` in comprehensions is as fast as a simple assignment
 ``y = expr``.

From ba401dc39a684b9ef8bb1118e68ace8f9458e2ce Mon Sep 17 00:00:00 2001
From: Serhiy Storchaka <storchaka@gmail.com>
Date: Fri, 10 Jan 2020 10:55:15 +0200
Subject: [PATCH 6/6] Add tests for unpacking a single star expression.

---
 Lib/test/test_dictcomps.py |  5 +++++
 Lib/test/test_genexps.py   |  7 +++++++
 Lib/test/test_listcomps.py |  7 +++++++
 Lib/test/test_setcomps.py  | 13 ++++++++++---
 4 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/Lib/test/test_dictcomps.py b/Lib/test/test_dictcomps.py
index 0fa1079085f673..16aa651b93c46b 100644
--- a/Lib/test/test_dictcomps.py
+++ b/Lib/test/test_dictcomps.py
@@ -122,6 +122,11 @@ def test_assignment_idiom_in_comprehensions(self):
         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 377b273316229d..86e4e195f55ec5 100644
--- a/Lib/test/test_genexps.py
+++ b/Lib/test/test_genexps.py
@@ -24,6 +24,13 @@
     >>> 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 da441da7d28b07..62b3319ad936d7 100644
--- a/Lib/test/test_listcomps.py
+++ b/Lib/test/test_listcomps.py
@@ -25,6 +25,13 @@
     >>> [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_setcomps.py b/Lib/test/test_setcomps.py
index ad954b04259b86..ecc4fffec0d849 100644
--- a/Lib/test/test_setcomps.py
+++ b/Lib/test/test_setcomps.py
@@ -23,13 +23,20 @@
 
 Test the idiom for temporary variable assignment in comprehensions.
 
-    >>> list(sorted({j*j for i in range(4) for j in [i+1]}))
+    >>> sorted({j*j for i in range(4) for j in [i+1]})
     [1, 4, 9, 16]
-    >>> list(sorted({j*k for i in range(4) for j in [i+1] for k in [j+1]}))
+    >>> sorted({j*k for i in range(4) for j in [i+1] for k in [j+1]})
     [2, 6, 12, 20]
-    >>> list(sorted({j*k for i in range(4) for j, k in [(i+1, i+2)]}))
+    >>> 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