[pypy-commit] pypy default: Attempting to add a JitDriver to unpackiterable(generator).

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: 
Changeset: r48642:894e6aa4bb16
Date: 2011-11-01 11:37 +0100
http://bitbucket.org/pypy/pypy/changeset/894e6aa4bb16/

Log:	Attempting to add a JitDriver to unpackiterable(generator).

diff --git a/pypy/interpreter/baseobjspace.py b/pypy/interpreter/baseobjspace.py
--- a/pypy/interpreter/baseobjspace.py
+++ b/pypy/interpreter/baseobjspace.py
@@ -778,6 +778,11 @@
         Raise an OperationError(w_ValueError) if the length is wrong."""
         w_iterator = self.iter(w_iterable)
         if expected_length == -1:
+            # xxx special hack for speed
+            from pypy.interpreter.generator import GeneratorIterator
+            if isinstance(w_iterator, GeneratorIterator):
+                return w_iterator.unpackiterable()
+            # /xxx
             return self._unpackiterable_unknown_length(w_iterator, w_iterable)
         else:
             return self._unpackiterable_known_length(w_iterator,
diff --git a/pypy/interpreter/generator.py b/pypy/interpreter/generator.py
--- a/pypy/interpreter/generator.py
+++ b/pypy/interpreter/generator.py
@@ -155,3 +155,32 @@
                                                  "interrupting generator of ")
                     break
                 block = block.previous
+
+    def unpackiterable(self):
+        """This is a hack for performance: runs the generator and collects
+        all produced items in a list."""
+        # XXX copied and simplified version of send_ex()
+        space = self.space
+        if self.running:
+            raise OperationError(space.w_ValueError,
+                                 space.wrap('generator already executing'))
+        results_w = []
+        frame = self.frame
+        if frame is None:    # already finished
+            return results_w
+        self.running = True
+        try:
+            while True:
+                jitdriver.jit_merge_point(frame=frame)
+                w_result = frame.execute_frame(space.w_None)
+                # if the frame is now marked as finished, it was RETURNed from
+                if frame.frame_finished_execution:
+                    break
+                results_w.append(w_result)     # YIELDed
+        finally:
+            frame.f_backref = jit.vref_None
+            self.running = False
+            self.frame = None
+        return results_w
+
+jitdriver = jit.JitDriver(greens=['frame.pycode'], reds=['frame'])
diff --git a/pypy/interpreter/test/test_generator.py b/pypy/interpreter/test/test_generator.py
--- a/pypy/interpreter/test/test_generator.py
+++ b/pypy/interpreter/test/test_generator.py
@@ -267,3 +267,9 @@
         assert r.startswith("<generator object myFunc at 0x")
         assert list(g) == [1]
         assert repr(g) == r
+
+    def test_unpackiterable_gen(self):
+        g = (i*i for i in range(-5, 3))
+        assert set(g) == set([0, 1, 4, 9, 16, 25])
+        assert set(g) == set()
+        assert set(i for i in range(0)) == set()