[pypy-commit] pypy default: Look inside RPython generators too.

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r52184:a1be520f19fd
Date: 2012-02-07 15:10 +0100
http://bitbucket.org/pypy/pypy/changeset/a1be520f19fd/

Log:	Look inside RPython generators too.

diff --git a/pypy/jit/codewriter/flatten.py b/pypy/jit/codewriter/flatten.py
--- a/pypy/jit/codewriter/flatten.py
+++ b/pypy/jit/codewriter/flatten.py
@@ -162,7 +162,9 @@
         if len(block.exits) == 1:
             # A single link, fall-through
             link = block.exits[0]
-            assert link.exitcase is None
+            assert link.exitcase in (None, False, True)
+            # the cases False or True should not really occur, but can show
+            # up in the manually hacked graphs for generators...
             self.make_link(link)
         #
         elif block.exitswitch is c_last_exception:
diff --git a/pypy/jit/metainterp/test/test_ajit.py b/pypy/jit/metainterp/test/test_ajit.py
--- a/pypy/jit/metainterp/test/test_ajit.py
+++ b/pypy/jit/metainterp/test/test_ajit.py
@@ -3706,6 +3706,18 @@
         # here it works again
         self.check_operations_history(guard_class=0, record_known_class=1)
 
+    def test_generator(self):
+        def g(n):
+            yield n+1
+            yield n+2
+            yield n+3
+        def f(n):
+            gen = g(n)
+            return gen.next() * gen.next() * gen.next()
+        res = self.interp_operations(f, [10])
+        assert res == 11 * 12 * 13
+        self.check_operations_history(int_add=3, int_mul=2)
+
 
 class TestLLtype(BaseLLtypeTests, LLJitMixin):
     def test_tagged(self):
diff --git a/pypy/translator/generator.py b/pypy/translator/generator.py
--- a/pypy/translator/generator.py
+++ b/pypy/translator/generator.py
@@ -68,6 +68,7 @@
         (next_entry, return_value) = func(entry)
         self.current = next_entry
         return return_value
+    next._jit_look_inside_ = True
     GeneratorIterator.next = next
     return func   # for debugging