"""Flow graph building for generators"""

from rpython.flowspace.argument import Signature

from rpython.flowspace.bytecode import HostCode

from rpython.flowspace.pygraph import PyGraph

from rpython.flowspace.model import (Block, Link, Variable,

Constant, checkgraph, const)

from rpython.flowspace.operation import op

from rpython.translator.unsimplify import insert_empty_startblock, split_block

from rpython.translator.simplify import eliminate_empty_blocks, simplify_graph

from rpython.tool.sourcetools import func_with_new_name

class AbstractPosition(object):

_immutable_ = True

_attrs_ = ()

def make_generator_entry_graph(func):

# This is the first copy of the graph. We replace it with

# a small bootstrap graph.

code = HostCode._from_code(func.func_code)

graph = PyGraph(func, code)

block = graph.startblock

for name, w_value in zip(code.co_varnames, block.framestate.mergeable):

if isinstance(w_value, Variable):

w_value.rename(name)

varnames = get_variable_names(graph.startblock.inputargs)

GeneratorIterator = make_generatoriterator_class(varnames)

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)

return graph

def tweak_generator_graph(graph):

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

class GeneratorIterator(object):

class Entry(AbstractPosition):

_immutable_ = True

varnames = var_names

def __init__(self, entry):

self.current = entry

def __iter__(self):

return self

return GeneratorIterator

def replace_graph_with_bootstrap(GeneratorIterator, graph):

Entry = GeneratorIterator.Entry

newblock = Block(graph.startblock.inputargs)

op_entry = op.simple_call(const(Entry))

v_entry = op_entry.result

newblock.operations.append(op_entry)

assert len(graph.startblock.inputargs) == len(Entry.varnames)

for v, name in zip(graph.startblock.inputargs, Entry.varnames):

newblock.operations.append(op.setattr(v_entry, Constant(name), v))

op_generator = op.simple_call(const(GeneratorIterator), v_entry)

newblock.operations.append(op_generator)

newblock.closeblock(Link([op_generator.result], graph.returnblock))

graph.startblock = newblock

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

assert entry is not None # else, recursive generator invocation

(next_entry, return_value) = func(entry)

self.current = next_entry

return return_value

GeneratorIterator.next = next

graph._tweaked_func = func # for testing

def get_variable_names(variables):

seen = set()

result = []

for v in variables:

name = v._name.strip('_')

while name in seen:

name += '_'

result.append('g_' + name)

seen.add(name)

return result

def _insert_reads(block, varnames):

assert len(varnames) == len(block.inputargs)

v_entry1 = Variable('entry')

for i, name in enumerate(varnames):

hlop = op.getattr(v_entry1, const(name))

hlop.result = block.inputargs[i]

block.operations.insert(i, hlop)

block.inputargs = [v_entry1]

def tweak_generator_body_graph(Entry, graph):

# First, always run simplify_graph in order to reduce the number of

# variables passed around

simplify_graph(graph)

insert_empty_startblock(graph)

_insert_reads(graph.startblock, Entry.varnames)

Entry.block = graph.startblock

#

mappings = [Entry]

#

stopblock = Block([])

op0 = op.simple_call(const(StopIteration))

op1 = op.type(op0.result)

stopblock.operations = [op0, op1]

stopblock.closeblock(Link([op1.result, op0.result], graph.exceptblock))

#

for block in list(graph.iterblocks()):

for exit in block.exits:

if exit.target is graph.returnblock:

exit.args = []

exit.target = stopblock

assert block is not stopblock

for index in range(len(block.operations)-1, -1, -1):

hlop = block.operations[index]

if hlop.opname == 'yield_':

[v_yielded_value] = hlop.args

del block.operations[index]

newlink = split_block(block, index)

newblock = newlink.target

varnames = get_variable_names(newlink.args)

#

class Resume(AbstractPosition):

_immutable_ = True

_attrs_ = varnames

block = newblock

Resume.__name__ = 'Resume%d' % len(mappings)

mappings.append(Resume)

#

_insert_reads(newblock, varnames)

#

op_resume = op.simple_call(const(Resume))

block.operations.append(op_resume)

v_resume = op_resume.result

for i, name in enumerate(varnames):

block.operations.append(

op.setattr(v_resume, const(name), newlink.args[i]))

op_pair = op.newtuple(v_resume, v_yielded_value)

block.operations.append(op_pair)

newlink.args = [op_pair.result]

newlink.target = graph.returnblock

#

regular_entry_block = Block([Variable('entry')])

block = regular_entry_block

for Resume in mappings:

op_check = op.isinstance(block.inputargs[0], const(Resume))

block.operations.append(op_check)

block.exitswitch = op_check.result

link1 = Link([block.inputargs[0]], Resume.block)

link1.exitcase = True

nextblock = Block([Variable('entry')])

link2 = Link([block.inputargs[0]], nextblock)

link2.exitcase = False

block.closeblock(link1, link2)

block = nextblock

block.closeblock(Link([Constant(AssertionError),

Constant(AssertionError("bad generator class"))],

graph.exceptblock))

graph.startblock = regular_entry_block

graph.signature = Signature(['entry'])

graph.defaults = ()

checkgraph(graph)

eliminate_empty_blocks(graph)