# -*- coding: utf-8 -*-

"""
Function inlining.
"""

from __future__ import print_function, division, absolute_import

from pykit.error import CompileError
from pykit.analysis import loop_detection
from pykit.ir import Function, Builder, findallops, copy_function, verify
from pykit.transform import ret as ret_normalization

def rewrite_return(func):
    """Rewrite ret ops to assign to a variable instead, which is returned"""
    ret_normalization.run(func)
    [ret] = findallops(func, 'ret')
    [value] = ret.args
    ret.delete()
    return value

def inline(func, call):
    """
    Inline the call instruction into func.

    :return: { old_op : new_op }
    """
    callee = call.args[0]
    # assert_inlinable(func, call, callee, uses)

    builder = Builder(func)
    builder.position_before(call)
    inline_header, inline_exit = builder.splitblock()
    new_callee, valuemap = copy_function(callee, temper=func.temp)
    result = rewrite_return(new_callee)

    # Fix up arguments
    for funcarg, arg in zip(new_callee.args, call.args[1]):
        funcarg.replace_uses(arg)

    # Copy blocks
    new_blocks = list(new_callee.blocks)
    after = inline_header
    for block in new_blocks:
        block.parent = None
        func.add_block(block, after=after)
        after = block

    # Fix up wiring
    builder.jump(new_callee.startblock)
    with builder.at_end(new_callee.exitblock):
        builder.jump(inline_exit)

    # Replicate exc_setup
    setups = [op for op in call.block.leaders if op.opcode == 'exc_setup']
    for block in new_blocks:
        builder.position_at_beginning(block)
        for setup in setups:
            builder.exc_setup(setup.args[0])

    # Fix up final result of call
    if result is not None:
        # non-void return
        result.unlink()
        result.result = call.result
        call.replace(result)
    else:
        call.delete()

    func.reset_uses()
    verify(func)

    return valuemap

def assert_inlinable(func, call, callee, uses):
    """
    Verify that a function call can be inlined.

    We can inline generators if they are consumed in a single loop:

        - iter(g) must be in a loop header
        - next(g) must be in the loop body

    :return: None if inlineable, or an exception with a message
    """
    if not isinstance(callee, Function):
        return CompileError("Cannot inline external function: %s" % (callee,))

    yields = findallops(callee, 'yield')
    if yields:
        for use in uses[call]:
            if use.opcode not in ('iter', 'next'):
                return CompileError(
                    "Cannot inline generator with use %s" % (use,))

        if len(uses[call]) != 2:
            return CompileError("Can only")
        loops = loop_detection.find_natural_loops(func)

@@ -691,16 +691,43 @@
 args%5B0%5D%0A
+    callblock = call.block%0A
     # as
@@ -767,16 +767,16 @@
 , uses)%0A
-
 %0A    bui
@@ -1468,251 +1468,62 @@
     
-# Replicate exc_setup%0A    setups = %5Bop for op in call.block.leaders if op.opcode == 'exc_setup'%5D%0A    for block in new_blocks:%0A        builder.position_at_beginning(block)%0A        for setup in setups:%0A            builder.exc_setup(setup.args%5B0%5D
+stretch_exception_block(builder, callblock, new_blocks
 )%0A%0A 
@@ -1770,16 +1770,16 @@
 (func)%0A%0A
-
     retu
@@ -1791,16 +1791,1313 @@
 luemap%0A%0A
+%0Adef stretch_exception_block(builder, originating_block, new_blocks):%0A    %22%22%22%0A    Replicate exc_setup across %60new_blocks%60 that were introduced from%0A    %60originating_block%60 by inlining.%0A    %22%22%22%0A    # Replicate exc_setup%0A    setups = %5Bop for op in originating_block.leaders if op.opcode == 'exc_setup'%5D%0A    for block in new_blocks:%0A        builder.position_at_beginning(block)%0A        for setup in setups:%0A            builder.exc_setup(setup.args%5B0%5D)%0A%0A    for setup in setups:%0A        %5Bsuccessors%5D = setup.args # exception handling blocks%0A%0A        # Patch phis in exception handling blocks for new predecessors%0A        for successor in successors:%0A            for op in successor.leaders:%0A                if op.opcode == 'phi':%0A                    blocks, values = map(list, op.args)%0A%0A                    # Find value for original predecessor%0A                    idx = blocks.index(originating_block)%0A                    value = values%5Bidx%5D%0A%0A                    # Patch value for new predecessors%0A                    for block in new_blocks:%0A                        if block not in blocks:%0A                            # We found a new predecessor!%0A                            blocks.append(block)%0A                            values.append(value)%0A%0A                    op.set_args(%5Bblocks, values%5D)%0A%0A%0A
 def asse

Patch phis in exception handler blocks after inlining