[pypy-commit] pypy generator-in-rpython: Generators in the flow space: keep the structure of the graph,

[arigo]

Author: Armin Rigo <arigo at tunes.org>
Branch: generator-in-rpython
Changeset: r49496:a084a39f04b9
Date: 2011-11-14 16:41 +0100
http://bitbucket.org/pypy/pypy/changeset/a084a39f04b9/

Log:	Generators in the flow space: keep the structure of the graph, but
	insert 'generator_entry' as the first operation and emit 'yield'
	operations.

diff --git a/pypy/objspace/flow/flowcontext.py b/pypy/objspace/flow/flowcontext.py
--- a/pypy/objspace/flow/flowcontext.py
+++ b/pypy/objspace/flow/flowcontext.py
@@ -185,7 +185,7 @@
 class FlowExecutionContext(ExecutionContext):
 
     def __init__(self, space, code, globals, constargs={}, outer_func=None,
-                 name=None):
+                 name=None, is_generator=False):
         ExecutionContext.__init__(self, space)
         self.code = code
 
@@ -208,6 +208,7 @@
         initialblock = SpamBlock(FrameState(frame).copy())
         self.pendingblocks = collections.deque([initialblock])
         self.graph = FunctionGraph(name or code.co_name, initialblock)
+        self.is_generator = is_generator
 
     make_link = Link # overridable for transition tracking
 
@@ -247,6 +248,8 @@
         return outcome, w_exc_cls, w_exc_value
 
     def build_flow(self):
+        if self.is_generator:
+            self.produce_generator_entry()
         while self.pendingblocks:
             block = self.pendingblocks.popleft()
             frame = self.create_frame()
@@ -259,9 +262,15 @@
                 self.topframeref = jit.non_virtual_ref(frame)
                 self.crnt_frame = frame
                 try:
-                    w_result = frame.dispatch(frame.pycode,
-                                              frame.last_instr,
-                                              self)
+                    frame.frame_finished_execution = False
+                    while True:
+                        w_result = frame.dispatch(frame.pycode,
+                                                  frame.last_instr,
+                                                  self)
+                        if frame.frame_finished_execution:
+                            break
+                        else:
+                            self.generate_yield(frame, w_result)
                 finally:
                     self.crnt_frame = None
                     self.topframeref = old_frameref
@@ -307,6 +316,19 @@
             del self.recorder
         self.fixeggblocks()
 
+    def produce_generator_entry(self):
+        [initialblock] = self.pendingblocks
+        initialblock.operations.append(
+            SpaceOperation('generator_entry', list(initialblock.inputargs),
+                           Variable()))
+
+    def generate_yield(self, frame, w_result):
+        assert self.is_generator
+        self.recorder.crnt_block.operations.append(
+            SpaceOperation('yield', [w_result], Variable()))
+        frame.pushvalue(None)
+        frame.last_instr += 1
+
     def fixeggblocks(self):
         # EggBlocks reuse the variables of their previous block,
         # which is deemed not acceptable for simplicity of the operations
diff --git a/pypy/objspace/flow/objspace.py b/pypy/objspace/flow/objspace.py
--- a/pypy/objspace/flow/objspace.py
+++ b/pypy/objspace/flow/objspace.py
@@ -8,6 +8,7 @@
 from pypy.interpreter.pycode import PyCode, cpython_code_signature
 from pypy.interpreter.module import Module
 from pypy.interpreter.error import OperationError
+from pypy.interpreter.astcompiler.consts import CO_GENERATOR
 from pypy.interpreter import pyframe, argument
 from pypy.objspace.flow.model import *
 from pypy.objspace.flow import flowcontext, operation, specialcase
@@ -247,9 +248,7 @@
         if func.func_doc and func.func_doc.lstrip().startswith('NOT_RPYTHON'):
             raise Exception, "%r is tagged as NOT_RPYTHON" % (func,)
         code = func.func_code
-        if code.co_flags & 32:
-            # generator
-            raise TypeError("%r is a generator" % (func,))
+        is_generator = bool(code.co_flags & CO_GENERATOR)
         code = PyCode._from_code(self, code)
         if func.func_closure is None:
             cl = None
@@ -265,7 +264,8 @@
         class outerfunc: # hack
             closure = cl
         ec = flowcontext.FlowExecutionContext(self, code, func.func_globals,
-                                              constargs, outerfunc, name)
+                                              constargs, outerfunc, name,
+                                              is_generator)
         graph = ec.graph
         graph.func = func
         # attach a signature and defaults to the graph
diff --git a/pypy/objspace/flow/test/test_generator.py b/pypy/objspace/flow/test/test_generator.py
new file mode 100644
--- /dev/null
+++ b/pypy/objspace/flow/test/test_generator.py
@@ -0,0 +1,18 @@
+from pypy.objspace.flow.test.test_objspace import Base
+
+
+class TestGenerator(Base):
+
+    def test_simple_generator(self):
+        def f(n):
+            i = 0
+            while i < n:
+                yield i
+                yield i
+                i += 1
+        graph = self.codetest(f)
+        ops = self.all_operations(graph)
+        assert ops == {'generator_entry': 1,
+                       'lt': 1, 'is_true': 1,
+                       'yield': 2,
+                       'inplace_add': 1}
diff --git a/pypy/objspace/flow/test/test_objspace.py b/pypy/objspace/flow/test/test_objspace.py
--- a/pypy/objspace/flow/test/test_objspace.py
+++ b/pypy/objspace/flow/test/test_objspace.py
@@ -882,12 +882,6 @@
                 num = bytecode_spec.opmap[name]
                 flow_meth_names[num] = locals()['old_' + name]
 
-    def test_generator(self):
-        def f():
-            yield 3
-
-        py.test.raises(TypeError, "self.codetest(f)")
-
     def test_dont_capture_RuntimeError(self):
         class Foo:
             def __hash__(self):