"""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
import modelverse_jit.runtime as jit_runtime
import modelverse_jit.bytecode_to_tree as bytecode_to_tree

# 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

    def lower(self, state):
        """Lowers this 'simple' block to a tree."""
        return tree_ir.create_block(
            state.lower_block(self.body),
            self.next_block.lower(state))

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

    def lower(self, state):
        """Lowers this 'loop' block to a tree."""
        old_loop = state.current_loop
        state.current_loop = self
        inner = tree_ir.LoopInstruction(
            tree_ir.create_block(
                self.inner.lower(state),
                tree_ir.BreakInstruction()))
        state.current_loop = old_loop
        return tree_ir.create_block(
            inner,
            self.next_block.lower(state))

class MultipleBlock(object):
    """A 'multiple' block in the relooper algorithm that does _not_ require a loop."""
    def __init__(self, handled_blocks, next_block):
        self.handled_blocks = handled_blocks
        self.next_block = next_block

    def lower_handled_blocks(self, state):
        """Lowers the handled blocks of this 'multiple' block to a tree."""
        result = tree_ir.EmptyInstruction()
        for entry, block in self.handled_blocks:
            result = tree_ir.SelectInstruction(
                tree_ir.BinaryInstruction(
                    state.label_variable,
                    '==',
                    tree_ir.LiteralInstruction(entry.index)),
                block.lower(state),
                result)
        return result

    def lower(self, state):
        """Lowers this 'multiple' block to a tree."""
        return tree_ir.create_block(
            self.lower_handled_blocks(state),
            self.next_block.lower(state))

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

    def lower(self, state):
        """Lowers this 'multiple' block to a tree."""
        old_loop = state.current_loop
        state.current_loop = self
        inner = tree_ir.LoopInstruction(
            tree_ir.create_block(
                self.lower_handled_blocks(state),
                tree_ir.BreakInstruction()))
        state.current_loop = old_loop
        return tree_ir.create_block(
            inner,
            self.next_block.lower(state))

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:
        if block in next_blocks:
            continue

        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)

    reachable_from_inner, inner_component = solipsize(
        graph_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 = MultipleLoopBlock({}, 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)

class LoweringState(object):
    """Stores information related to the relooper->tree conversion."""
    def __init__(self, jit):
        self.jit = jit
        self.label_variable = tree_ir.VariableName('__label')
        self.definition_loads = {}
        self.local_name_map = bytecode_to_tree.LocalNameMap()
        self.root_edge_names = {}

    def __get_root_edge_name(self, root_node):
        """Gets the name of the given root edge's variable."""
        return self.__get_root_node_name(root_node) + '_edge'

    def __get_root_node_name(self, root_node):
        """Gets the name of the given root node's variable."""
        if isinstance(root_node, cfg_ir.Definition):
            return self.__get_root_edge_name(root_node.value)

        if root_node in self.root_edge_names:
            return self.root_edge_names[root_node]

        result = 'jit_locals%d' % len(self.root_edge_names)
        self.root_edge_names[root_node] = result
        return result

    def __create_value_load(self, value):
        """Creates a tree that loads the given value."""
        if value.has_value():
            if isinstance(value, cfg_ir.Literal):
                return self.lower_literal(value)
            else:
                return tree_ir.LoadLocalInstruction(None)
        else:
            return tree_ir.LiteralInstruction(None)

    def load_definition(self, definition):
        """Loads the given definition's variable."""
        if definition in self.definition_loads:
            return self.definition_loads[definition]

        result = self.__create_value_load(definition.value)
        self.definition_loads[definition] = result
        return result

    def lower_block(self, block):
        """Lowers the given (relooped) block to a tree."""
        statements = []
        for definition in block.definitions:
            statements.append(self.lower_definition(definition))
        statements.append(self.lower_flow(block.flow))
        return tree_ir.create_block(statements)

    def lower_definition(self, definition):
        """Lowers the given definition to a tree."""
        instruction = definition.value
        tree_instruction = self.lower_value(instruction)
        def_load = self.load_definition(definition)
        if isinstance(def_load, tree_ir.LocalInstruction):
            return def_load.create_store(tree_instruction)
        else:
            return tree_instruction

    def lower_value(self, value):
        """Lowers the given instruction to a tree."""
        value_type = type(value)
        if value_type in LoweringState.value_lowerings:
            return LoweringState.value_lowerings[value_type](self, value)
        else:
            raise jit_runtime.JitCompilationFailedException(
                "Unknown CFG instruction: '%s'" % value)

    def lower_literal(self, value):
        """Lowers the given literal value."""
        return tree_ir.LiteralInstruction(value.literal)

    def lower_check_local_exists(self, value):
        """Lowers a 'check-value-exists' value."""
        return tree_ir.LocalExistsInstruction(
            self.local_name_map.get_local_name(value.variable.node_id))

    def lower_declare_local(self, value):
        """Lowers a 'declare-local' value."""
        local_name = self.local_name_map.get_local_name(value.variable.node_id)
        return tree_ir.create_block(
            tree_ir.create_new_local_node(
                local_name,
                self.load_definition(value.root_node)),
            tree_ir.LoadLocalInstruction(local_name))

    def lower_resolve_local(self, value):
        """Lowers a 'resolve-local' value."""
        return tree_ir.LoadLocalInstruction(
            self.local_name_map.get_local_name(value.variable.node_id))

    def lower_alloc_root_node(self, value):
        """Lowers an 'alloc-root-node' value."""
        local_name = tree_ir.VariableName(self.__get_root_node_name(value))
        return tree_ir.create_block(
            tree_ir.create_new_local_node(
                local_name,
                tree_ir.LoadIndexInstruction(
                    tree_ir.LoadLocalInstruction(jit_runtime.KWARGS_PARAMETER_NAME),
                    tree_ir.LiteralInstruction('task_root')),
                self.__get_root_edge_name(value)),
            tree_ir.LoadLocalInstruction(local_name))

    def lower_free_root_node(self, value):
        """Lowers a 'free-root-node' value."""
        return tree_ir.DeleteEdgeInstruction(self.__get_root_edge_name(value.root_node))

    def lower_load_pointer(self, value):
        """Lowers a 'load' value."""
        return bytecode_to_tree.create_access(self.load_definition(value.pointer))

    def lower_store_pointer(self, value):
        """Lowers a 'store' value."""
        return bytecode_to_tree.create_assign(
            self.load_definition(value.pointer), self.load_definition(value.value))

    value_lowerings = {
        cfg_ir.Literal : lower_literal,
        cfg_ir.CheckLocalExists : lower_check_local_exists,
        cfg_ir.DeclareLocal : lower_declare_local,
        cfg_ir.ResolveLocal : lower_resolve_local,
        cfg_ir.AllocateRootNode : lower_alloc_root_node,
        cfg_ir.DeallocateRootNode : lower_free_root_node,
        cfg_ir.LoadPointer : lower_load_pointer,
        cfg_ir.StoreAtPointer : lower_store_pointer
    }

    def lower_flow(self, flow):
        """Lowers the given (relooped) flow instruction to a tree."""
        flow_type = type(flow)
        if flow_type in LoweringState.flow_lowerings:
            return LoweringState.flow_lowerings[flow_type](self, flow)
        else:
            raise jit_runtime.JitCompilationFailedException(
                "Unknown CFG flow instruction: '%s'" % flow)

    def lower_jump(self, flow):
        """Lowers the given 'jump' flow instruction to a tree."""
        return self.lower_branch(flow.branch)

    def lower_select(self, flow):
        """Lowers the given 'select' flow instruction to a tree."""
        return tree_ir.SelectInstruction(
            self.load_definition(flow.condition),
            self.lower_branch(flow.if_branch),
            self.lower_branch(flow.else_branch))

    def lower_return(self, flow):
        """Lowers the given 'return' flow instruction to a tree."""
        return tree_ir.ReturnInstruction(self.load_definition(flow.value))

    def lower_unreachable(self, _):
        """Lowers the given 'unreachable' flow instruction to a tree."""
        return tree_ir.EmptyInstruction()

    def lower_break(self, _):
        """Lowers the given 'break' flow instruction to a tree."""
        return tree_ir.BreakInstruction()

    def lower_continue(self, _):
        """Lowers the given 'continue' flow instruction to a tree."""
        return tree_ir.ContinueInstruction()

    def lower_branch(self, branch):
        """Lowers the given (relooped) branch to a tree."""
        for param, arg in zip(branch.block.parameters, branch.arguments):
            self.load_definition(param).create_store(self.load_definition(arg))

        return tree_ir.StoreLocalInstruction(
            self.label_variable,
            tree_ir.LiteralInstruction(branch.block.index))

    flow_lowerings = {
        cfg_ir.JumpFlow : lower_jump,
        cfg_ir.SelectFlow : lower_select,
        cfg_ir.ReturnFlow : lower_return,
        cfg_ir.UnreachableFlow : lower_unreachable,
        BreakFlow : lower_break,
        ContinueFlow : lower_continue
    }