[pypy-commit] pypy default: Remove the need for the RPython manual optimization of writing "x in (2,

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r56715:6363bce133c9
Date: 2012-08-12 23:30 +0200
http://bitbucket.org/pypy/pypy/changeset/6363bce133c9/

Log:	Remove the need for the RPython manual optimization of writing "x in
	(2, 3, 4, 5)" because it is more efficient than "x in [2, 3, 4, 5]".

diff --git a/pypy/rpython/test/test_rlist.py b/pypy/rpython/test/test_rlist.py
--- a/pypy/rpython/test/test_rlist.py
+++ b/pypy/rpython/test/test_rlist.py
@@ -686,6 +686,31 @@
                 res = self.interpret(fn, [i, case])
                 assert res is fn(i, case)
 
+    def test_constant_list_contains(self):
+        # a 'contains' operation on list containing only annotation-time
+        # constants should be optimized into the equivalent code of
+        # 'in prebuilt-dictionary'.  Hard to test directly...
+        def g():
+            return 16
+        def f(i):
+            return i in [1, 2, 4, 8, g()]
+        res = self.interpret(f, [2])
+        assert res is True
+        res = self.interpret(f, [15])
+        assert res is False
+        res = self.interpret(f, [16])
+        assert res is True
+
+    def test_nonconstant_list_contains(self):
+        def f(i):
+            return i in [1, -i, 2, 4, 8]
+        res = self.interpret(f, [2])
+        assert res is True
+        res = self.interpret(f, [15])
+        assert res is False
+        res = self.interpret(f, [0])
+        assert res is True
+
 
     def test_not_a_char_list_after_all(self):
         def fn():
diff --git a/pypy/translator/transform.py b/pypy/translator/transform.py
--- a/pypy/translator/transform.py
+++ b/pypy/translator/transform.py
@@ -109,6 +109,32 @@
                                         op.result)
                 block.operations[i] = new_op
 
+# x in [2, 3]
+# -->
+# b = newlist(2, 3)
+# c = contains(b, x)
+# -->
+# c = contains(Constant((2, 3)), x)
+
+def transform_list_contains(self, block_subset):
+    """Transforms x in [2, 3]"""
+    for block in block_subset:
+        newlist_sources = {}    # maps b to [2, 3] in the above notation
+        for i in range(len(block.operations)):
+            op = block.operations[i]
+            if op.opname == 'newlist':
+                newlist_sources[op.result] = op.args
+            elif op.opname == 'contains' and op.args[0] in newlist_sources:
+                items = {}
+                for v in newlist_sources[op.args[0]]:
+                    s = self.binding(v)
+                    if not s.is_immutable_constant():
+                        break
+                    items[s.const] = None
+                else:
+                    # all arguments of the newlist are annotation constants
+                    op.args[0] = Constant(items)
+
 
 def transform_dead_op_vars(self, block_subset):
     # we redo the same simplification from simplify.py,
@@ -221,6 +247,7 @@
     transform_allocate,
     transform_extend_with_str_slice,
     transform_extend_with_char_count,
+    transform_list_contains,
     ]
 
 def transform_graph(ann, extra_passes=None, block_subset=None):