"""Converts 'declare-local', 'load' and 'store' instructions into SSA form."""

from collections import defaultdict
import modelverse_jit.cfg_ir as cfg_ir

def get_local_id(def_or_value):
    """Gets the node of the local resolved or declared by the given definition or value.
       If the given definition or value does not refer to a 'resolve-local' or
       'declare-local' node, then None is returned."""
    value = cfg_ir.get_def_value(def_or_value)
    if isinstance(value, (cfg_ir.ResolveLocal, cfg_ir.DeclareLocal)):
        return value.variable.node_id
    else:
        return None

def get_ineligible_local_ids(entry_point):
    """Finds the ids of all local variables which are not eligible for conversion to SSA form."""
    # Local variables are eligible for conversion to SSA form if their pointer node is never
    # leaked to the outside world. So we know that we can safely convert a local to SSA form
    # if 'resolve-local' values are only used by 'load' and 'store' values.
    ineligible_local_ids = set()
    def __maybe_mark_ineligible(def_or_value):
        local_id = get_local_id(def_or_value)
        if local_id is not None:
            ineligible_local_ids.add(local_id)

    for block in cfg_ir.get_all_blocks(entry_point):
        for definition in block.definitions + [block.flow]:
            value = cfg_ir.get_def_value(definition)
            if isinstance(value, cfg_ir.LoadPointer):
                # Loading a pointer to a local is fine.
                pass
            elif isinstance(value, cfg_ir.StoreAtPointer):
                # Storing a value in a local is fine, too.
                # But be careful not to ignore a store where the stored value is a local pointer.
                __maybe_mark_ineligible(value.value)
            else:
                # Walk over the dependencies, and mark them all as ineligible for
                # local-to-SSA conversion.
                for dependency in value.get_all_dependencies():
                    __maybe_mark_ineligible(dependency)

    return ineligible_local_ids

def construct_ssa_form(entry_point):
    """Converts local variables into SSA form in the graph defined by the given entry point.
       A potentially new entry block is returned."""
    # Build some helper data structures.
    all_blocks = list(cfg_ir.get_all_blocks(entry_point))
    ineligible_locals = get_ineligible_local_ids(entry_point)
    predecessor_map = cfg_ir.get_all_predecessor_blocks(entry_point)

    # Create the SSA construction state.
    state = SSAConstructionState(all_blocks, ineligible_locals, predecessor_map)

    # Fill all blocks in the graph.
    for block in all_blocks:
        state.fill_block(block)

    # Update branches.
    for block in all_blocks:
        state.update_block_branches(block)

    # Nullify entry point parameters.
    return nullify_entry_block_parameters(entry_point)

def nullify_entry_block_parameters(entry_point):
    """Creates or returns an entry block that takes no parameters.
       The (new) entry block is returned."""
    # The SSA construction algorithm has the nasty habit of replacing potentially
    # undefined variables by entry block parameters. Codegen doesn't like that one
    # bit: it assumes that all block parameters are defined by the branches to those
    # parameters -- but there usually is no branch to the entry point!
    #
    # We can fix this by either replacing the block-parameters in the entry block by
    # definitions, or by creating a new entry point that jumps to the old entry
    # point with an argument list. The former construction is more efficient, but it is only
    # correct if there are no pre-existing branches to the entry point.
    #
    # We could build the predecessor map block and use the first option if possible,
    # but we won't; trivial phi elimination followed by block merging will reduce
    # the second construction to the first one anyway.
    if len(entry_point.parameters) == 0:
        return entry_point

    pre_entry_point = cfg_ir.BasicBlock(entry_point.counter)
    arg_list = []
    for _ in entry_point.parameters:
        literal_def = pre_entry_point.append_definition(cfg_ir.Literal(None))
        arg_list.append(literal_def)

    pre_entry_point.flow = cfg_ir.create_jump(entry_point, arg_list)
    return pre_entry_point

# The algorithms below are based on
# Simple and Efficient Construction of Static Single Assignment Form by M. Braun et al
# (https://pp.info.uni-karlsruhe.de/uploads/publikationen/braun13cc.pdf).

class SSAConstructionState(object):
    """Encapsulates state related to SSA construction."""
    def __init__(self, all_blocks, ineligible_locals, predecessor_map):
        self.all_blocks = all_blocks
        self.ineligible_locals = ineligible_locals
        self.predecessor_map = predecessor_map
        # `current_defs` is a local node id -> basic block -> definition map.
        self.current_defs = defaultdict(dict)
        # `incomplete_phis` is a basic block -> local node id -> block parameter def map.
        self.incomplete_phis = defaultdict(dict)
        # `extra_phi_operands` is a basic block -> block parameter def -> def map.
        self.extra_phi_operands = defaultdict(dict)
        self.processed_blocks = set()
        self.filled_blocks = set()
        self.sealed_blocks = set()

    def read_variable(self, block, node_id):
        """Reads the latest definition of the local variable with the
           given node id for the specified block."""
        if block in self.current_defs[node_id]:
            return self.current_defs[node_id][block]
        else:
            return self.read_variable_recursive(block, node_id)

    def write_variable(self, block, node_id, value):
        """Writes the given value to the local with the specified id in the
           specified block."""
        self.current_defs[node_id][block] = value

    def read_variable_recursive(self, block, node_id):
        """Reads the latest definition of the local variable with the
           given node id from one of the given block's predecessor blocks."""
        if block not in self.sealed_blocks:
            # Create an incomplete phi.
            val = block.append_parameter(cfg_ir.BlockParameter())
            self.incomplete_phis[block][node_id] = val
        elif len(self.predecessor_map[block]) == 1:
            # Optimize the common case of one predecessor: no phi needed.
            pred = next(iter(self.predecessor_map[block]))
            val = self.read_variable(pred, node_id)
        else:
            # Break potential cycles with an operandless phi.
            val = block.append_parameter(cfg_ir.BlockParameter())
            self.write_variable(block, node_id, val)
            val = self.add_phi_operands(node_id, val)

        self.write_variable(block, node_id, val)
        return val

    def add_phi_operands(self, node_id, phi_def):
        """Finds out which arguments branches should provide for the given block
           parameter definition."""
        # Determine operands from predecessors
        all_values = []
        for pred in self.predecessor_map[phi_def.block]:
            arg = self.read_variable(pred, node_id)
            self.extra_phi_operands[pred][phi_def] = arg
            all_values.append(arg)
        return self.try_remove_trivial_phi(phi_def, all_values)

    def try_remove_trivial_phi(self, phi_def, values):
        """Tries to remove a trivial block parameter definition."""
        # This is a somewhat simplified (and less powerful) version of the
        # algorithm in the SSA construction paper. That's kind of okay, though;
        # trivial phi elimination is also implemented as a separate pass in the
        # optimization pipeline.
        trivial_phi_val = cfg_ir.get_trivial_phi_value(phi_def, values)
        if trivial_phi_val is None:
            return phi_def
        else:
            for pred in self.predecessor_map[phi_def.block]:
                del self.extra_phi_operands[pred][phi_def]

            phi_def.block.remove_parameter(phi_def)
            phi_def.redefine(trivial_phi_val)
            phi_def.block.prepend_definition(phi_def)
            return trivial_phi_val

    def has_sealed(self, block):
        """Tells if the given block has been sealed yet."""
        return block in self.sealed_blocks

    def can_seal(self, block):
        """Tells if the given block can be sealed right away."""
        # A block can be sealed if all if its predecessors have been filled.
        return all(
            [predecessor in self.filled_blocks for predecessor in self.predecessor_map[block]])

    def seal_all_sealable_blocks(self):
        """Seals all sealable blocks."""
        for block in self.all_blocks:
            if self.can_seal(block):
                self.seal_block(block)

    def seal_block(self, block):
        """Seals the given block."""
        if self.has_sealed(block):
            return

        for node_id, phi_def in self.incomplete_phis[block].items():
            self.add_phi_operands(node_id, phi_def)

        self.sealed_blocks.add(block)

    def has_filled(self, block):
        """Tells if the given block has been filled yet."""
        return block in self.filled_blocks

    def fill_block(self, block):
        """Visits all definitions in the given block. Locals are converted into SSA form."""
        if block in self.processed_blocks:
            return

        self.processed_blocks.add(block)

        # Try to seal the block right away if at all possible.
        if self.can_seal(block):
            self.seal_block(block)

        block_definitions = list(block.definitions)
        for definition in block_definitions:
            value = definition.value
            if cfg_ir.is_value_def(value, cfg_ir.LoadPointer):
                # Read the variable from the definitions dictionary.
                node_id = get_local_id(value.pointer)
                if node_id is not None and node_id not in self.ineligible_locals:
                    definition.redefine(self.read_variable(block, node_id))
            elif isinstance(value, cfg_ir.StoreAtPointer):
                node_id = get_local_id(value.pointer)
                if node_id is not None and node_id not in self.ineligible_locals:
                    # Write to the variable, and replace the definition by a 'None' literal.
                    self.write_variable(block, node_id, value.value)
                    definition.redefine(cfg_ir.Literal(None))
            elif isinstance(value, cfg_ir.DeclareLocal):
                node_id = value.variable.node_id
                if node_id not in self.ineligible_locals:
                    definition.redefine(cfg_ir.Literal(None))


        # Mark the block as filled.
        self.filled_blocks.add(block)

        # Seal all sealable blocks.
        self.seal_all_sealable_blocks()

        # Fill successor blocks.
        for branch in block.flow.branches():
            self.fill_block(branch.block)

    def update_block_branches(self, block):
        """Appends arguments to the given block's flow instruction's branches, if necessary."""
        for branch in block.flow.branches():
            # Find all pairs phis which are defined in the branch target block.
            applicable_pairs = [
                (phi_def, operand_def)
                for phi_def, operand_def in self.extra_phi_operands[block].items()
                if phi_def.block == branch.block]

            if len(applicable_pairs) == 0:
                # We might as well early-out here.
                continue

            # Sort the pairs by block parameter index.
            sorted_pairs = sorted(
                applicable_pairs,
                key=lambda (phi_def, _): phi_def.block.parameters.index(phi_def))

            # Append arguments to the branch.
            for _, arg in sorted_pairs:
                branch.arguments.append(arg)