"""Lowers CFG-IR to tree-IR."""

import modelverse_jit.cfg_ir as cfg_ir
import modelverse_jit.cfg_optimization as cfg_optimization
import modelverse_jit.tree_ir as tree_ir

# The CFG deconstruction code here is based on the relooper algorithm
# as detailed in https://github.com/kripken/Relooper/blob/master/paper.pdf

class FlowGraphComponent(object):
    """Represents a control-flow graph component."""
    def __init__(self, entry_blocks, blocks, reachable):
        self.entry_blocks = entry_blocks
        self.blocks = blocks
        self.reachable = reachable

    def get_entry_reachable_blocks(self):
        """Gets the set of all blocks that are reachable from the entry points."""
        return set.union(*[self.get_reachable_blocks(block) for block in self.entry_blocks])

    def get_reachable_blocks(self, block):
        """Gets all blocks in this component that are reachable from the given block."""
        return set.intersection(self.reachable[block], self.blocks)

    def get_directly_reachable_blocks(self, block):
        """Gets all blocks in this component that are reachable from the given block by taking
           exactly one branch."""
        return set.intersection(cfg_optimization.get_directly_reachable_blocks(block), self.blocks)

    def can_reach(self, source_block, target_block):
        """Tests if the first block can reach the second."""
        return target_block in self.reachable[source_block] and target_block in self.blocks

    def sub_component(self, new_entry_points):
        """Creates a sub-component of this component, with the given new entry points."""
        result = FlowGraphComponent(new_entry_points, self.blocks, self.reachable)
        result.blocks = result.get_entry_reachable_blocks()
        return result

class SimpleBlock(object):
    """A 'simple' block in the relooper algorithm."""
    def __init__(self, body, next_block):
        self.body = body
        self.next_block = next_block

class LoopBlock(object):
    """A 'loop' block in the relooper algorithm."""
    def __init__(self, inner, next_block):
        self.inner = inner
        self.next_block = next_block

class MultipleBlock(object):
    """A 'multiple' block in the relooper algorithm."""
    def __init__(self, handled_blocks, next_block):
        self.handled_blocks = handled_blocks
        self.next_block = next_block

class ContinueFlow(cfg_ir.FlowInstruction):
    """Represents a control flow instruction which continues to the next loop iteration."""
    def __init__(self, loop):
        cfg_ir.FlowInstruction.__init__(self)
        self.loop = loop

    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        return []

    def branches(self):
        """Gets a list of basic blocks targeted by this flow instruction."""
        return []

    def __str__(self):
        return 'continue'

class BreakFlow(cfg_ir.FlowInstruction):
    """Represents a control flow instruction which breaks out of a loop."""
    def __init__(self, loop):
        cfg_ir.FlowInstruction.__init__(self)
        self.loop = loop

    def get_dependencies(self):
        """Gets all definitions and instructions on which this instruction depends."""
        return []

    def branches(self):
        """Gets a list of basic blocks targeted by this flow instruction."""
        return []

    def __str__(self):
        return 'break'

def solipsize(graph_component, loop, entry_blocks, next_blocks):
    """Replaces branches to entry blocks and next blocks by jumps to 'continue' and 'break'
       blocks, respectively."""
    all_blocks = set()
    reachable_from_inner = set()
    for block in graph_component.blocks:
        all_blocks.add(block)
        for branch in block.flow.branches():
            if branch.block in entry_blocks:
                # Create a 'continue' block, and point to that.
                continue_block = cfg_ir.BasicBlock(block.counter)
                for param in branch.block.parameters:
                    continue_block.append_parameter(param)

                continue_block.flow = ContinueFlow(loop)
                branch.block = continue_block
                all_blocks.add(continue_block)
            elif branch.block in next_blocks:
                # Create a 'break' block, and point to that.
                break_block = cfg_ir.BasicBlock(block.counter)
                for param in branch.block.parameters:
                    break_block.append_parameter(param)

                break_block.flow = BreakFlow(loop)
                branch.block = break_block
                all_blocks.add(break_block)
                reachable_from_inner.add(branch.block)

    return reachable_from_inner, FlowGraphComponent(
        graph_component.entry_blocks,
        all_blocks,
        {
            block : cfg_optimization.get_reachable_blocks(block)
            for block in all_blocks
        })

def to_relooper_loop(graph_component):
    """Converts the given graph component to a relooper 'loop'."""
    entry_blocks = graph_component.entry_blocks
    result = LoopBlock(None, None)
    inner_blocks = []
    next_blocks = []
    for block in graph_component.blocks:
        if any([graph_component.can_reach(block, ep) for ep in entry_blocks]):
            inner_blocks.append(block)
        else:
            next_blocks.append(block)

    inner_component = FlowGraphComponent(
        entry_blocks, inner_blocks, graph_component.reachable)
    reachable_from_inner, inner_component = solipsize(
        inner_component, result, entry_blocks, next_blocks)
    result.inner = reloop(inner_component)

    next_component = FlowGraphComponent(
        reachable_from_inner, next_blocks, graph_component.reachable)
    result.next_block = reloop(next_component)

    return result

def to_relooper_multiple_or_loop(graph_component):
    """Converts the given graph component to a relooper 'multiple' or 'loop'."""
    # From the Emscripten paper:
    #
    # If we have more than one entry, try to create a Multiple block:
    # For each entry, find all the labels it reaches that
    # cannot be reached by any other entry. If at least one entry
    # has such labels, return a Multiple block, whose Handled
    # blocks are blocks for those labels (and whose entries are
    # those labels), and whose Next block is all the rest. Entries
    # for the next block are entries that did not become part of
    # the Handled blocks, and also labels that can be reached
    # from the Handled blocks.
    entry_blocks = graph_component.entry_blocks
    if len(entry_blocks) <= 1:
        return to_relooper_loop(graph_component)

    entry_reachables = {ep : graph_component.get_reachable_blocks(ep) for ep in entry_blocks}
    exclusive_entries = {}
    for entry in entry_blocks:
        exclusive_blocks = set(entry_reachables[entry])
        for other_entry in entry_blocks:
            if other_entry != entry:
                exclusive_blocks.difference_update(entry_reachables[other_entry])
        if len(exclusive_blocks) > 0:
            exclusive_entries[entry] = exclusive_blocks

    if len(exclusive_entries) == 0:
        return to_relooper_loop(graph_component)

    next_entries = set(graph_component.entry_blocks)
    for block_set in exclusive_entries.values():
        for elem in block_set:
            directly_reachable = graph_component.get_directly_reachable_blocks(elem)
            directly_reachable.remove(elem)
            next_entries.update(directly_reachable)

    next_entries.difference_update(exclusive_entries.keys())

    result = MultipleBlock({}, None)
    for entry, exclusive_blocks in exclusive_entries.items():
        other_blocks = set(graph_component.blocks)
        other_blocks.difference_update(exclusive_blocks)
        result.handled_blocks[entry] = reloop(
            solipsize(graph_component, result, set([entry]), other_blocks))

    result.next_block = reloop(graph_component.sub_component(next_entries))

    return result

def reloop(graph_component):
    """Applies the relooper algorithm to the given graph component."""
    reachable_set = graph_component.get_entry_reachable_blocks()
    entry_blocks = graph_component.entry_blocks
    if len(entry_blocks) == 1 and entry_blocks[0] not in reachable_set:
        graph_component.blocks.remove(entry_blocks[0])
        return SimpleBlock(
            entry_blocks[0],
            reloop(
                FlowGraphComponent(
                    graph_component.get_directly_reachable_blocks(entry_blocks[0]),
                    graph_component.blocks,
                    graph_component.reachable)))
    elif all([block in reachable_set for block in entry_blocks]):
        return to_relooper_loop(graph_component)
    else:
        return to_relooper_multiple_or_loop(graph_component)