Commits

Armin Rigo committed a1be520

Look inside RPython generators too.

Comments (0)

Files changed (3)

pypy/jit/codewriter/flatten.py

         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:

pypy/jit/metainterp/test/test_ajit.py

         # 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):

pypy/translator/generator.py

         (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