[pypy-commit] pypy py3.6: do the same optimization that CPython does for (1, 2, 3, *a)

[cfbolz]

Author: Carl Friedrich Bolz-Tereick <cfbolz at gmx.de>
Branch: py3.6
Changeset: r96368:e41b9e56b6cc
Date: 2019-03-27 22:35 +0100
http://bitbucket.org/pypy/pypy/changeset/e41b9e56b6cc/

Log:	do the same optimization that CPython does for (1, 2, 3, *a) (but on
	the AST level)

diff --git a/pypy/interpreter/astcompiler/optimize.py b/pypy/interpreter/astcompiler/optimize.py
--- a/pypy/interpreter/astcompiler/optimize.py
+++ b/pypy/interpreter/astcompiler/optimize.py
@@ -302,6 +302,9 @@
                 node = tup.elts[i]
                 w_const = node.as_constant()
                 if w_const is None:
+                    new_elts = self._optimize_constant_star_unpacks(tup.elts)
+                    if new_elts is not None:
+                        return ast.Tuple(new_elts, ast.Load, tup.lineno, tup.col_offset)
                     return tup
                 consts_w[i] = w_const
             # intern the string constants packed into the tuple here,
@@ -314,6 +317,62 @@
         w_consts = self.space.newtuple(consts_w)
         return ast.Constant(w_consts, tup.lineno, tup.col_offset)
 
+    def _make_starred_tuple_const(self, consts_w, firstelt):
+        w_consts = self.space.newtuple(consts_w)
+        return ast.Starred(ast.Constant(
+                    w_consts, firstelt.lineno, firstelt.col_offset),
+                ast.Load, firstelt.lineno, firstelt.col_offset)
+
+    def _optimize_constant_star_unpacks(self, elts):
+        # turn (1, 2, 3, *a) into (*(1, 2, 3), *a) with a constant (1, 2, 3)
+        # or similarly, for lists
+        contains_starred = False
+        for i in range(len(elts)):
+            elt = elts[i]
+            if isinstance(elt, ast.Starred):
+                contains_starred = True
+                break
+        if not contains_starred:
+            return None
+        new_elts = []
+        changed = False
+        const_since_last_star_w = []
+        after_last_star_index = 0
+        for i in range(len(elts)):
+            elt = elts[i]
+            if isinstance(elt, ast.Starred):
+                if const_since_last_star_w is not None:
+                    firstelt = elts[after_last_star_index]
+                    new_elts.append(self._make_starred_tuple_const(
+                        const_since_last_star_w, firstelt))
+                    changed = True
+                const_since_last_star_w = []
+                after_last_star_index = i + 1
+                new_elts.append(elt)
+            elif const_since_last_star_w is not None:
+                w_const = elt.as_constant()
+                if w_const is None:
+                    new_elts.extend(elts[after_last_star_index:i + 1])
+                    const_since_last_star_w = None
+                else:
+                    const_since_last_star_w.append(w_const)
+            else:
+                new_elts.append(elt)
+        if after_last_star_index != len(elts) and const_since_last_star_w is not None:
+            firstelt = elts[after_last_star_index]
+            new_elts.append(self._make_starred_tuple_const(
+                const_since_last_star_w, firstelt))
+            changed = True
+        if changed:
+            return new_elts
+
+    def visit_List(self, l):
+        if l.ctx == ast.Load and l.elts:
+            new_elts = self._optimize_constant_star_unpacks(l.elts)
+            if new_elts:
+                return ast.List(new_elts, ast.Load, l.lineno, l.col_offset)
+        return l
+
     def visit_Subscript(self, subs):
         if subs.ctx == ast.Load:
             w_obj = subs.value.as_constant()
diff --git a/pypy/interpreter/astcompiler/test/test_compiler.py b/pypy/interpreter/astcompiler/test/test_compiler.py
--- a/pypy/interpreter/astcompiler/test/test_compiler.py
+++ b/pypy/interpreter/astcompiler/test/test_compiler.py
@@ -1099,14 +1099,23 @@
                 return a, b, c
         """
         yield self.st, func, "f()", (1, [2, 3], 4)
+
+    def test_unpacking_while_building(self):
         func = """def f():
             b = [4,5,6]
-            c = 7
-            a = [*b, c]
+            a = (*b, 7)
+            return a
+        """
+        yield self.st, func, "f()", (4, 5, 6, 7)
+
+        func = """def f():
+            b = [4,5,6]
+            a = [*b, 7]
             return a
         """
         yield self.st, func, "f()", [4, 5, 6, 7]
 
+
     def test_extended_unpacking_fail(self):
         exc = py.test.raises(SyntaxError, self.simple_test, "*a, *b = [1, 2]",
                              None, None).value
@@ -1542,6 +1551,31 @@
         counts = self.count_instructions(source)
         assert ops.BUILD_TUPLE not in counts
 
+    def test_constant_tuples_star(self):
+        source = """def f(a, c):
+            return (u"a", 1, *a, 3, 5, 3, *c)
+        """
+        counts = self.count_instructions(source)
+        assert ops.BUILD_TUPLE not in counts
+
+        source = """def f(a, c, d):
+            return (u"a", 1, *a, c, 1, *d, 1, 2, 3)
+        """
+        counts = self.count_instructions(source)
+        assert counts[ops.BUILD_TUPLE] == 1
+
+    def test_constant_list_star(self):
+        source = """def f(a, c):
+            return [u"a", 1, *a, 3, 5, 3, *c]
+        """
+        counts = self.count_instructions(source)
+        assert ops.BUILD_TUPLE not in counts
+
+        source = """def f(a, c, d):
+            return [u"a", 1, *a, c, 1, *d, 1, 2, 3]
+        """
+        counts = self.count_instructions(source)
+        assert counts[ops.BUILD_TUPLE] == 1
 
 
 class TestHugeStackDepths: