[pypy-commit] pypy default: fold lists of constants into a constant tuple in the AST

[pjenvey]

Author: Philip Jenvey <pjenvey at underboss.org>
Branch: 
Changeset: r64298:a9007a4abb78
Date: 2013-05-18 16:37 -0700
http://bitbucket.org/pypy/pypy/changeset/a9007a4abb78/

Log:	fold lists of constants into a constant tuple in the AST

diff --git a/pypy/interpreter/astcompiler/codegen.py b/pypy/interpreter/astcompiler/codegen.py
--- a/pypy/interpreter/astcompiler/codegen.py
+++ b/pypy/interpreter/astcompiler/codegen.py
@@ -853,9 +853,10 @@
             self.emit_jump(ops.JUMP_IF_FALSE_OR_POP, cleanup, True)
             if i < (ops_count - 1):
                 comp.comparators[i].walkabout(self)
-        comp.comparators[-1].walkabout(self)
-        last_kind = compare_operations(comp.ops[-1])
-        self.emit_op_arg(ops.COMPARE_OP, last_kind)
+        last_op, last_comparator = comp.ops[-1], comp.comparators[-1]
+        if not self._optimize_comparator(last_op, last_comparator):
+            last_comparator.walkabout(self)
+        self.emit_op_arg(ops.COMPARE_OP, compare_operations(last_op))
         if ops_count > 1:
             end = self.new_block()
             self.emit_jump(ops.JUMP_FORWARD, end)
@@ -864,6 +865,30 @@
             self.emit_op(ops.POP_TOP)
             self.use_next_block(end)
 
+    def _optimize_comparator(self, op, node):
+        """Fold lists of constants in the context of "in"/"not in".
+
+        lists are folded into tuples, otherwise returns False
+        """
+        if op in (ast.In, ast.NotIn) and isinstance(node, ast.List):
+            w_const = self._tuple_of_consts(node.elts)
+            if w_const is not None:
+                self.load_const(w_const)
+                return True
+        return False
+
+    def _tuple_of_consts(self, elts):
+        """Return a tuple of consts from elts if possible, or None"""
+        count = len(elts) if elts is not None else 0
+        consts_w = [None] * count
+        for i in range(count):
+            w_value = elts[i].as_constant()
+            if w_value is None:
+                # Not all constants
+                return None
+            consts_w[i] = w_value
+        return self.space.newtuple(consts_w)
+
     def visit_IfExp(self, ifexp):
         self.update_position(ifexp.lineno)
         end = self.new_block()
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
@@ -973,3 +973,16 @@
         counts = self.count_instructions(source3)
         assert counts[ops.BUILD_LIST] == 1
         assert ops.BUILD_LIST_FROM_ARG not in counts
+
+    def test_folding_of_list_constants(self):
+        for source in (
+            # in/not in constants with BUILD_LIST should be folded to a tuple:
+            'a in [1,2,3]',
+            'a not in ["a","b","c"]',
+            'a in [None, 1, None]',
+            'a not in [(1, 2), 3, 4]',
+            ):
+            source = 'def f(): %s' % source
+            counts = self.count_instructions(source)
+            assert ops.BUILD_LIST not in counts
+            assert ops.LOAD_CONST in counts
diff --git a/pypy/interpreter/test/test_compiler.py b/pypy/interpreter/test/test_compiler.py
--- a/pypy/interpreter/test/test_compiler.py
+++ b/pypy/interpreter/test/test_compiler.py
@@ -936,6 +936,13 @@
         output = s.getvalue()
         assert "LOAD_GLOBAL" not in output
 
+    def test_folding_of_list_constants(self):
+        source = 'a in [1, 2, 3]'
+        co = compile(source, '', 'exec')
+        i = co.co_consts.index((1, 2, 3))
+        assert i > -1
+        assert isinstance(co.co_consts[i], tuple)
+
 class AppTestCallMethod(object):
     spaceconfig = {'objspace.opcodes.CALL_METHOD': True}