Commits

Armin Rigo committed a4906ce

Finish the replacement of the flow graph of generators. The behavior of
the tweaked flow graphs is not tested so far; needs some RPython tests.

Comments (0)

Files changed (5)

pypy/objspace/flow/objspace.py

                 return ecls
         return None
 
-    def build_flow(self, func, constargs={}):
+    def build_flow(self, func, constargs={}, tweak_for_generator=True):
         """
         """
         if func.func_doc and func.func_doc.lstrip().startswith('NOT_RPYTHON'):
             e = error.FlowingError(formated)
             raise error.FlowingError, e, tb
         checkgraph(graph)
+        #
+        if is_generator and tweak_for_generator:
+            from pypy.translator.generator import tweak_generator_graph
+            tweak_generator_graph(graph)
+        #
         return graph
 
     def fixedview(self, w_tuple, expected_length=None):

pypy/objspace/flow/test/test_generator.py

                 yield i
                 yield i
                 i += 1
-        graph = self.codetest(f)
+        graph = self.codetest(f, tweak_for_generator=False)
         ops = self.all_operations(graph)
         assert ops == {'generator_mark': 1,
                        'lt': 1, 'is_true': 1,

pypy/objspace/flow/test/test_objspace.py

 is_operator = getattr(operator, 'is_', operator.eq) # it's not there 2.2
 
 class Base:
-    def codetest(self, func):
+    def codetest(self, func, **kwds):
         import inspect
         try:
             func = func.im_func
         except AttributeError:
             pass
         #name = func.func_name
-        graph = self.space.build_flow(func)
+        graph = self.space.build_flow(func, **kwds)
         graph.source = inspect.getsource(func)
         self.show(graph)
         return graph

pypy/translator/generator.py

 from pypy.translator.unsimplify import insert_empty_startblock
 from pypy.translator.unsimplify import split_block
 from pypy.translator.simplify import eliminate_empty_blocks
+from pypy.tool.sourcetools import func_with_new_name
 
 
 class AbstractPosition(object):
     _attrs_ = ()
 
 
-def replace_graph_with_bootstrap(graph, graph_of_body, Entry):
-    #
+def tweak_generator_graph(graph):
+    if not hasattr(graph.func, '_generator_next_method_of_'):
+        # This is the first copy of the graph.  We replace it with
+        # a small bootstrap graph.
+        GeneratorIterator = make_generatoriterator_class(graph)
+        replace_graph_with_bootstrap(GeneratorIterator, graph)
+        # We attach a 'next' method to the GeneratorIterator class
+        # that will invoke the real function, based on a second
+        # copy of the graph.
+        attach_next_method(GeneratorIterator, graph)
+    else:
+        # This is the second copy of the graph.  Tweak it.
+        GeneratorIterator = graph.func._generator_next_method_of_
+        tweak_generator_body_graph(GeneratorIterator.Entry, graph)
+
+
+def make_generatoriterator_class(graph):
     class GeneratorIterator(object):
-        graph = graph_of_body
+        class Entry(AbstractPosition):
+            varnames = get_variable_names(graph.startblock.inputargs)
         def __init__(self, entry):
             self.current = entry
-    GeneratorIterator.Entry = Entry
-    #
+    return GeneratorIterator
+
+def replace_graph_with_bootstrap(GeneratorIterator, graph):
+    Entry = GeneratorIterator.Entry
     newblock = Block(graph.startblock.inputargs)
     v_generator = Variable('generator')
     v_entry = Variable('entry')
                        v_generator))
     newblock.closeblock(Link([v_generator], graph.returnblock))
     graph.startblock = newblock
-    return GeneratorIterator
+
+def attach_next_method(GeneratorIterator, graph):
+    func = graph.func
+    func = func_with_new_name(func, '%s__next' % (func.func_name,))
+    func._generator_next_method_of_ = GeneratorIterator
+    func._always_inline_ = True
+    #
+    def next(self):
+        entry = self.current
+        self.current = None
+        (next_entry, return_value) = func(entry)
+        self.current = next_entry
+        return return_value
+    GeneratorIterator.next = next
+    return func   # for debugging
 
 def get_variable_names(variables):
     seen = set()
                            block.inputargs[i]))
     block.inputargs = [v_entry1]
 
-def tweak_generator_body_graph(graph):
+def tweak_generator_body_graph(Entry, graph):
     assert graph.startblock.operations[0].opname == 'generator_mark'
     graph.startblock.operations.pop(0)
     #
-    entryvarnames = get_variable_names(graph.startblock.inputargs)
     insert_empty_startblock(None, graph)
-    _insert_reads(graph.startblock, entryvarnames)
+    _insert_reads(graph.startblock, Entry.varnames)
+    Entry.block = graph.startblock
     #
-    class Entry(AbstractPosition):
-        block = graph.startblock
-        varnames = entryvarnames
     mappings = [Entry]
     #
     for block in list(graph.iterblocks()):
     graph.startblock = regular_entry_block
     checkgraph(graph)
     eliminate_empty_blocks(graph)
-    try:
-        graph.func._always_inline_ = True
-    except AttributeError:
-        pass
-    return Entry

pypy/translator/test/test_generator.py

 from pypy.objspace.flow.objspace import FlowObjSpace
 from pypy.objspace.flow.model import Variable
 from pypy.translator.translator import TranslationContext
+from pypy.translator.generator import make_generatoriterator_class
 from pypy.translator.generator import replace_graph_with_bootstrap
 from pypy.translator.generator import get_variable_names
 from pypy.translator.generator import tweak_generator_body_graph
+from pypy.translator.generator import attach_next_method
+from pypy.translator.simplify import join_blocks
 
 
 # ____________________________________________________________
             yield n
         #
         space = FlowObjSpace()
-        graph = space.build_flow(func)
+        graph = space.build_flow(func, tweak_for_generator=False)
         assert graph.startblock.operations[0].opname == 'generator_mark'
-        class Entry:
-            varnames = ['g_n', 'g_x', 'g_y', 'g_z']
-        replace_graph_with_bootstrap(graph, 'newgraph', Entry)
+        GeneratorIterator = make_generatoriterator_class(graph)
+        replace_graph_with_bootstrap(GeneratorIterator, graph)
         if option.view:
             graph.show()
         block = graph.startblock
             z -= 10
         #
         space = FlowObjSpace()
-        graph = space.build_flow(f)
-        tweak_generator_body_graph(graph)
+        graph = space.build_flow(f, tweak_for_generator=False)
+        class Entry:
+            varnames = ['g_n', 'g_x', 'g_y', 'g_z']
+        tweak_generator_body_graph(Entry, graph)
         if option.view:
             graph.show()
         # XXX how to test directly that the graph is correct?  :-(
+
+    def test_tweak_generator_graph(self):
+        def f(n, x, y, z):
+            z *= 10
+            yield n + 1
+            z -= 10
+        #
+        space = FlowObjSpace()
+        graph = space.build_flow(f, tweak_for_generator=False)
+        GeneratorIterator = make_generatoriterator_class(graph)
+        replace_graph_with_bootstrap(GeneratorIterator, graph)
+        func1 = attach_next_method(GeneratorIterator, graph)
+        if option.view:
+            graph.show()
+        #
+        assert func1._generator_next_method_of_ is GeneratorIterator
+        assert hasattr(GeneratorIterator, 'next')
+        #
+        graph_next = space.build_flow(GeneratorIterator.next.im_func)
+        join_blocks(graph_next)
+        if option.view:
+            graph_next.show()
+        #
+        graph1 = space.build_flow(func1, tweak_for_generator=False)
+        tweak_generator_body_graph(GeneratorIterator.Entry, graph1)
+        if option.view:
+            graph1.show()
+
+    def test_automatic(self):
+        def f(n, x, y, z):
+            z *= 10
+            yield n + 1
+            z -= 10
+        #
+        space = FlowObjSpace()
+        graph = space.build_flow(f)   # tweak_for_generator=True
+        if option.view:
+            graph.show()
+        block = graph.startblock
+        assert len(block.exits) == 1
+        assert block.exits[0].target is graph.returnblock