peephole: fix UNPACK_SEQUENCE optim

vstinner · vstinner · commit 53f91add60d7 · 2016-03-02T02:17:12.000+01:00
diff --git a/bytecode/peephole_opt.py b/bytecode/peephole_opt.py
@@ -193,27 +193,45 @@ def replace_container_of_consts(self, instr, container_type):
         self.replace_load_const(instr.arg, instr, value)
 
     def build_tuple_unpack_seq(self, instr):
-        nconst = instr.arg
-        if nconst < 1:
-            return
-
         next_instr = self.get_next_instr('UNPACK_SEQUENCE')
         if next_instr is None or next_instr.arg != instr.arg:
             return
 
-        start = self.index - 1
+        if instr.arg < 1:
+            return
 
-        # Rewrite LOAD_CONST instructions in the reverse order
-        load_consts = self.block[start - nconst:start]
-        self.block[start - nconst:start] = reversed(load_consts)
+        if (self.const_stack
+            and instr.arg <= len(self.const_stack)):
+            nconst = instr.arg
+            start = self.index - 1
 
-        # Remove BUILD_TUPLE+UNPACK_SEQUENCE
-        self.block[start:start + 2] = ()
-        self.index -= 2
-        self.const_stack.clear()
+            # Rewrite LOAD_CONST instructions in the reverse order
+            load_consts = self.block[start - nconst:start]
+            self.block[start - nconst:start] = reversed(load_consts)
 
-        # FIXME: why not rewriting LOAD_CONST in the reverse order to support
-        # any number of aguments, rather than using ROT_TWO/ROT_THREE tricks?
+            # Remove BUILD_TUPLE+UNPACK_SEQUENCE
+            self.block[start:start + 2] = ()
+            self.index -= 2
+            self.const_stack.clear()
+            return
+
+        if instr.arg == 1:
+            # Replace BUILD_TUPLE 1 + UNPACK_SEQUENCE 1 with NOP
+            del self.block[self.index-1:self.index+1]
+        elif instr.arg == 2:
+            # Replace BUILD_TUPLE 2 + UNPACK_SEQUENCE 2 with ROT_TWO
+            rot2 = Instr('ROT_TWO', lineno=instr.lineno)
+            self.block[self.index - 1:self.index+1] = (rot2,)
+            self.index -= 1
+            self.const_stack.clear()
+        elif instr.arg == 3:
+            # Replace BUILD_TUPLE 3 + UNPACK_SEQUENCE 3
+            # with ROT_THREE + ROT_TWO
+            rot3 = Instr('ROT_THREE', lineno=instr.lineno)
+            rot2 = Instr('ROT_TWO', lineno=instr.lineno)
+            self.block[self.index-1:self.index+1] = (rot3, rot2)
+            self.index -= 1
+            self.const_stack.clear()
 
     def build_tuple(self, instr, container_type):
         if instr.arg > len(self.const_stack):
diff --git a/bytecode/tests/test_peephole_opt.py b/bytecode/tests/test_peephole_opt.py
@@ -243,8 +243,9 @@ def test_build_list_unpack_seq(self):
                              Instr('STORE_NAME', 'x'),
                              Instr('STORE_NAME', 'y')])
             self.check(code,
-                       Instr('LOAD_NAME', 'b'),
                        Instr('LOAD_NAME', 'a'),
+                       Instr('LOAD_NAME', 'b'),
+                       Instr('ROT_TWO'),
                        Instr('STORE_NAME', 'x'),
                        Instr('STORE_NAME', 'y'))
 
@@ -258,13 +259,47 @@ def test_build_list_unpack_seq(self):
                              Instr('STORE_NAME', 'y'),
                              Instr('STORE_NAME', 'z')])
             self.check(code,
-                       Instr('LOAD_NAME', 'c'),
-                       Instr('LOAD_NAME', 'b'),
                        Instr('LOAD_NAME', 'a'),
+                       Instr('LOAD_NAME', 'b'),
+                       Instr('LOAD_NAME', 'c'),
+                       Instr('ROT_THREE'),
+                       Instr('ROT_TWO'),
                        Instr('STORE_NAME', 'x'),
                        Instr('STORE_NAME', 'y'),
                        Instr('STORE_NAME', 'z'))
 
+    def test_build_tuple_unpack_seq_const(self):
+        # x, y = (3, 4)
+        code = Bytecode([Instr('LOAD_CONST', 3),
+                         Instr('LOAD_CONST', 4),
+                         Instr('BUILD_TUPLE', 2),
+                         Instr('UNPACK_SEQUENCE', 2),
+                         Instr('STORE_NAME', 'x'),
+                         Instr('STORE_NAME', 'y')])
+        self.check(code,
+                   Instr('LOAD_CONST', (3, 4)),
+                   Instr('UNPACK_SEQUENCE', 2),
+                   Instr('STORE_NAME', 'x'),
+                   Instr('STORE_NAME', 'y'))
+
+    def test_build_list_unpack_seq_const(self):
+        # x, y, z = [3, 4, 5]
+        code = Bytecode([Instr('LOAD_CONST', 3),
+                         Instr('LOAD_CONST', 4),
+                         Instr('LOAD_CONST', 5),
+                         Instr('BUILD_LIST', 3),
+                         Instr('UNPACK_SEQUENCE', 3),
+                         Instr('STORE_NAME', 'x'),
+                         Instr('STORE_NAME', 'y'),
+                         Instr('STORE_NAME', 'z')])
+        self.check(code,
+                   Instr('LOAD_CONST', 5),
+                   Instr('LOAD_CONST', 4),
+                   Instr('LOAD_CONST', 3),
+                   Instr('STORE_NAME', 'x'),
+                   Instr('STORE_NAME', 'y'),
+                   Instr('STORE_NAME', 'z'))
+
     def test_build_set(self):
         # test = x in {1, 2, 3}
         code = Bytecode([Instr('LOAD_NAME', 'x'),
diff --git a/doc/peephole.rst b/doc/peephole.rst
@@ -89,8 +89,10 @@ Optimizations implemented in the peephole optimizer:
     * a << b, a >> b, a & b, a | b, a ^ b
 
   - BUILD_TUPLE: convert to a constant
-  - BUILD_TUPLE <n> + UNPACK_SEQUENCE <n> and BUILD_LIST <n>
-    + UNPACK_SEQUENCE <n>: rewrite LOAD_CONST instructions in the reverse order
+  - Replace BUILD_TUPLE <n> + UNPACK_SEQUENCE <n> and BUILD_LIST <n>
+    + UNPACK_SEQUENCE <n> with ROT_TWO (2 arguments) or ROT_THREE+ROT_TWO (3
+    arguments). For BUILD_LIST, if inputs are LOAD_CONST, rewrite LOAD_CONST in
+    the reverse order.
   - BUILD_LIST + COMPARE_OP(in/not in): convert list to tuple
   - BUILD_SET + COMPARE_OP(in/not in): convert set to frozenset
   - COMPARE_OP:
