"""Optimizes and analyzes CFG-IR."""

from collections import defaultdict
import modelverse_jit.cfg_ir as cfg_ir
import modelverse_jit.cfg_dominators as cfg_dominators
import modelverse_jit.cfg_ssa_construction as cfg_ssa_construction
import modelverse_jit.cfg_data_structures as cfg_data_structures
import modelverse_kernel.primitives as primitive_functions

def is_empty_block(block):
    """Tests if the given block contains no parameters or definitions."""
    return len(block.parameters) == 0 and len(block.definitions) == 0

def optimize_flow(block):
    """Optimizes the given block's flow instruction."""
    changed = True
    while changed:
        changed = False

        # Select flow with a literal condition can be optimized to a direct jump.
        if (isinstance(block.flow, cfg_ir.SelectFlow)
                and cfg_ir.is_literal_def(block.flow.condition)):
            literal = cfg_ir.get_literal_def_value(block.flow.condition)
            block.flow = cfg_ir.JumpFlow(
                block.flow.if_branch if literal else block.flow.else_branch)
            changed = True

        # Jumps to blocks which contain no parameters or definitions can be replaced
        # by the target block's flow.
        if (isinstance(block.flow, cfg_ir.JumpFlow)
                and is_empty_block(block.flow.branch.block)
                and block.flow.branch.block is not block):
            block.flow = block.flow.branch.block.flow
            changed = True

        # Branches to blocks which contain nothing but a jump can be replaced by branches
        # to the jump's target.
        for branch in block.flow.branches():
            if (is_empty_block(branch.block)
                    and branch.block is not block
                    and isinstance(branch.block.flow, cfg_ir.JumpFlow)):
                new_branch = branch.block.flow.branch
                branch.block = new_branch.block
                branch.arguments = new_branch.arguments
                changed = True

def optimize_graph_flow(entry_point):
    """Optimizes all flow instructions in the graph defined by the given entry point."""
    for block in cfg_ir.get_all_blocks(entry_point):
        optimize_flow(block)

def merge_blocks(entry_point):
    """Merges blocks which have exactly one predecessor with said predecessor, if the
       predecessor has a jump flow instruction."""
    predecessor_map = cfg_ir.get_all_predecessor_blocks(entry_point)
    queue = set(predecessor_map.keys())
    queue.add(entry_point)
    def __do_merge(source, target):
        target_params = list(target.parameters)
        branch_args = list(source.flow.branch.arguments)
        for target_param, branch_arg in zip(target_params, branch_args):
            target.remove_parameter(target_param)
            target_param.redefine(branch_arg)
            source.append_definition(target_param)

        target_defs = list(target.definitions)
        for target_def in target_defs:
            target.remove_definition(target_def)
            source.append_definition(target_def)

        source.flow = target.flow
        for preds in predecessor_map.values():
            if target in preds:
                preds[preds.index(target)] = source
                # preds.remove(target)
                # preds.add(source)

    while len(queue) > 0:
        block = queue.pop()
        if isinstance(block.flow, cfg_ir.JumpFlow):
            next_block = block.flow.branch.block
            preds = predecessor_map[next_block]
            if (len(preds) == 1
                    and next(iter(preds)) == block
                    and block != next_block
                    and next_block != entry_point):
                __do_merge(block, next_block)
                del predecessor_map[next_block]
                queue.add(block)
                if next_block in queue:
                    queue.remove(next_block)

def elide_local_checks(entry_point):
    """Tries to elide redundant checks on local variables."""
    # The plan here is to replace all check-local-exists defs by literals if
    # they are either dominated by an appropriate declare-local or not reachable
    # from a declare-local.
    local_checks = []
    local_defs = defaultdict(list)
    for block in cfg_ir.get_all_blocks(entry_point):
        for definition in block.definitions:
            def_value = definition.value
            if isinstance(def_value, cfg_ir.CheckLocalExists):
                local_checks.append((def_value.variable.node_id, definition))
            elif isinstance(def_value, cfg_ir.DeclareLocal):
                local_defs[def_value.variable.node_id].append(definition)

    dominator_tree = cfg_dominators.get_dominator_tree(entry_point)
    reachable_blocks = cfg_ir.get_all_reachable_blocks(entry_point)
    for (variable, check) in local_checks:
        is_reachable = False
        for local_def in local_defs[variable]:
            if dominator_tree.dominates_instruction(local_def, check):
                # Check is dominated by a definition. Replace it by a 'True' literal.
                check.redefine(cfg_ir.Literal(True))
                is_reachable = True
                break
            elif check.block in reachable_blocks[local_def.block]:
                is_reachable = True

        if not is_reachable:
            # Check cannot be reached from any definition. Replace it by a 'False' literal.
            check.redefine(cfg_ir.Literal(False))

def eliminate_unused_definitions(entry_point):
    """Tries to eliminate unused definitions in the control-flow graphb defined by the
       given entry point."""
    def_dependencies = defaultdict(set)
    root_defs = set()
    # Gather dependencies.
    for block in cfg_ir.get_all_blocks(entry_point):
        for definition in block.parameters + block.definitions:
            all_dependencies = list(definition.get_all_dependencies())
            def_dependencies[definition].update(
                [dep for dep in all_dependencies
                 if isinstance(dep, cfg_ir.Definition)])

            if len(all_dependencies) > 0 and definition.has_bidirectional_dependencies():
                for dep in all_dependencies:
                    def_dependencies[dep].add(definition)

            if definition.has_side_effects():
                root_defs.add(definition)

        for dep in block.flow.get_dependencies():
            if isinstance(dep, cfg_ir.Definition):
                root_defs.add(dep)
            else:
                assert isinstance(dep, cfg_ir.Branch)
                for param, arg in zip(dep.block.parameters, dep.arguments):
                    def_dependencies[param].add(arg)

    # Figure out which definitions are live.
    live_defs = set()
    def __mark_live(definition):
        if definition in live_defs:
            return

        live_defs.add(definition)
        if definition in def_dependencies:
            for dep in def_dependencies[definition]:
                __mark_live(dep)

    for root in root_defs:
        __mark_live(root)

    # Remove all dead definitions.
    dead_defs = set.difference(set(def_dependencies.keys()), live_defs)
    dead_phis = set()
    for dead_def in dead_defs:
        if isinstance(dead_def.value, cfg_ir.BlockParameter):
            dead_phis.add(dead_def)
        else:
            dead_def.block.remove_definition(dead_def)

    erase_parameters(entry_point, dead_phis)

def eliminate_trivial_phis(entry_point):
    """Eliminates trivial block parameters, i.e., block parameters which are really
       aliases."""
    phi_values = defaultdict(set)
    all_blocks = list(cfg_ir.get_all_blocks(entry_point))
    for block in all_blocks:
        for branch in block.flow.branches():
            for phi, arg in zip(branch.block.parameters, branch.arguments):
                phi_values[phi].add(arg)

    replacements = []
    for block in all_blocks:
        block_parameters = list(block.parameters)
        for parameter_def in block_parameters:
            trivial_phi_val = cfg_ir.get_trivial_phi_value(
                parameter_def, phi_values[parameter_def])
            if trivial_phi_val is not None:
                replacements.append((parameter_def, trivial_phi_val))

    erase_parameters(entry_point, set([parameter_def for parameter_def, _ in replacements]))

    for parameter_def, trivial_phi_val in replacements:
        block = parameter_def.block
        parameter_def.redefine(trivial_phi_val)
        block.prepend_definition(parameter_def)

def erase_parameters(entry_point, parameters_to_erase):
    """Erases all arguments for the given set of parameters, and then takes out the
       parameters themselves."""
    for block in cfg_ir.get_all_blocks(entry_point):
        for branch in block.flow.branches():
            new_arg_list = []
            for parameter, arg in zip(branch.block.parameters, branch.arguments):
                if parameter not in parameters_to_erase:
                    new_arg_list.append(arg)
            branch.arguments = new_arg_list

    for parameter_def in parameters_to_erase:
        parameter_def.block.remove_parameter(parameter_def)

def apply_cfg_intrinsic(intrinsic_function, original_definition, named_args):
    """Applies the given intrinsic to the given sequence of named arguments."""
    kwargs = dict(named_args)
    kwargs['original_def'] = original_definition
    return intrinsic_function(**kwargs)

def try_redefine_as_direct_call(definition, jit, called_globals):
    """Tries to redefine the given indirect call definition as a direct call."""
    call = cfg_ir.get_def_value(definition)
    if not isinstance(call, cfg_ir.IndirectFunctionCall):
        return

    target = cfg_ir.get_def_value(call.target)
    if isinstance(target, cfg_ir.LoadPointer):
        loaded_ptr = cfg_ir.get_def_value(target.pointer)
        if isinstance(loaded_ptr, cfg_ir.ResolveGlobal):
            resolved_var_name = loaded_ptr.variable.name
            called_globals.add(loaded_ptr)

            # Try to resolve the callee as an intrinsic.
            intrinsic = jit.get_cfg_intrinsic(resolved_var_name)
            if intrinsic is not None:
                definition.redefine(
                    cfg_ir.DirectFunctionCall(
                        resolved_var_name, call.argument_list,
                        cfg_ir.JIT_CFG_INTRINSIC_CALLING_CONVENTION))
            else:
                # Otherwise, build a thunk.
                thunk_name = jit.jit_thunk_global(resolved_var_name)
                calling_convention = (
                    cfg_ir.JIT_NO_GC_CALLING_CONVENTION
                    if jit.get_intrinsic(thunk_name) is not None
                    else cfg_ir.JIT_CALLING_CONVENTION)
                definition.redefine(
                    cfg_ir.DirectFunctionCall(
                        thunk_name, call.argument_list, calling_convention))
                called_globals.add(loaded_ptr)
    elif isinstance(target, cfg_ir.Literal):
        node_id = target.literal
        thunk_name = jit.jit_thunk_constant_function(node_id)
        definition.redefine(
            cfg_ir.DirectFunctionCall(
                thunk_name, call.argument_list, cfg_ir.JIT_CALLING_CONVENTION))

def get_checked_global(definition):
    """If the definition is a check that tests if a global does not exist, then
       the instruction that resolves the global is returned; otherwise None."""
    def_value = cfg_ir.get_def_value(definition)
    if not isinstance(def_value, cfg_ir.Binary):
        return None

    if def_value.operator != 'is':
        return None

    def __get_checked_global_single_dir(lhs, rhs):
        if (isinstance(lhs, cfg_ir.ResolveGlobal)
                and isinstance(rhs, cfg_ir.Literal)
                and rhs.literal is None):
            return lhs
        else:
            return None

    bin_lhs = cfg_ir.get_def_value(def_value.lhs)
    bin_rhs = cfg_ir.get_def_value(def_value.rhs)
    result = __get_checked_global_single_dir(bin_lhs, bin_rhs)
    if result is None:
        result = __get_checked_global_single_dir(bin_rhs, bin_lhs)

    return result

def optimize_calls(entry_point, jit):
    """Converts indirect calls to direct calls in the control-flow graph defined by the
       given entry point."""
    called_globals = set()
    global_exists_defs = defaultdict(list)
    for block in cfg_ir.get_all_blocks(entry_point):
        for definition in block.definitions:
            checked_global = get_checked_global(definition)
            if checked_global is not None:
                global_exists_defs[checked_global].append(definition)
            else:
                try_redefine_as_direct_call(definition, jit, called_globals)

    for resolve_global in called_globals:
        for exists_def in global_exists_defs[resolve_global]:
            exists_def.redefine(cfg_ir.Literal(False))

def expand_cfg_intrinsics(entry_point, jit):
    """Expands CFG JIT intrinsics in the control-flow graph defined by the given entry point."""
    for block in cfg_ir.get_all_blocks(entry_point):
        for definition in block.definitions:
            def_value = definition.value
            if (isinstance(def_value, cfg_ir.DirectFunctionCall)
                    and def_value.calling_convention ==
                    cfg_ir.JIT_CFG_INTRINSIC_CALLING_CONVENTION):
                intrinsic = jit.get_cfg_intrinsic(def_value.target_name)
                apply_cfg_intrinsic(intrinsic, definition, def_value.argument_list)

def simplify_values(entry_point):
    """Simplifies values in the control-flow graph defined by the given entry point."""
    for block in cfg_ir.get_all_blocks(entry_point):
        for definition in block.definitions:
            def_val = cfg_ir.get_def_value(definition)
            if isinstance(def_val, cfg_ir.Read):
                read_node = cfg_ir.get_def_value(def_val.node)
                if isinstance(read_node, cfg_ir.CreateNode):
                    definition.redefine(read_node.value)
            elif isinstance(def_val, cfg_ir.Binary):
                lhs = cfg_ir.get_def_value(def_val.lhs)
                rhs = cfg_ir.get_def_value(def_val.rhs)
                if isinstance(lhs, cfg_ir.Literal) and isinstance(rhs, cfg_ir.Literal):
                    definition.redefine(
                        cfg_ir.Literal(
                            eval('%r %s %r' % (lhs.literal, def_val.operator, rhs.literal))))
            elif isinstance(def_val, cfg_ir.Unary):
                operand = cfg_ir.get_def_value(def_val.operand)
                if isinstance(operand, cfg_ir.Literal):
                    definition.redefine(
                        cfg_ir.Literal(
                            eval('%s %r' % (def_val.operator, operand.literal))))

def inline_constants(entry_point):
    """Replaces reads of constant nodes by the literals they contain."""
    for block in cfg_ir.get_all_blocks(entry_point):
        for definition in block.definitions:
            def_val = cfg_ir.get_def_value(definition)
            if isinstance(def_val, cfg_ir.Read):
                read_node = cfg_ir.get_def_value(def_val.node)
                if isinstance(read_node, cfg_ir.Literal):
                    val, = yield [("RV", [read_node.literal])]
                    definition.redefine(cfg_ir.Literal(val))

def expand_indirect_definitions(entry_point):
    """Replaces indirect definitions by the values referred to by those definitions."""
    def __expand_indirect_defs(value):
        dependencies = value.get_dependencies()
        if len(dependencies) == 0:
            return value
        else:
            new_dependencies = []
            for dep in dependencies:
                new_dep = dep
                if isinstance(new_dep, cfg_ir.Definition):
                    while isinstance(new_dep.value, cfg_ir.Definition):
                        new_dep = new_dep.value
                else:
                    new_dep = __expand_indirect_defs(new_dep)

                new_dependencies.append(new_dep)
            return value.create(new_dependencies)

    for block in cfg_ir.get_all_blocks(entry_point):
        block_definitions = list(block.definitions)
        for definition in block_definitions:
            if isinstance(definition.value, cfg_ir.Definition):
                block.remove_definition(definition)
            else:
                definition.redefine(
                    __expand_indirect_defs(definition.value))

        block.flow = __expand_indirect_defs(block.flow)

def optimize_reads(entry_point):
    """Tries to replace repeated reads by a single read."""
    cfg_ir.match_and_rewrite(
        entry_point,
        lambda _: True,
        lambda use_def, _: cfg_ir.is_value_def(use_def, cfg_ir.Read),
        lambda def_def:
        def_def.redefine(
            cfg_ir.Read(def_def.insert_before(def_def.value))),
        lambda use_def, def_def: use_def.redefine(def_def))

def protect_from_gc(entry_point):
    """Protects locals in the control-flow graph defined by the given
       entry point from the GC."""
    root_node = entry_point.prepend_definition(cfg_ir.AllocateRootNode())
    def protect_def_from_gc(definition):
        """Protects the given definition from the GC."""
        definition.insert_after(cfg_ir.create_gc_protect(definition, root_node))

    def maybe_protect_def_from_gc(definition):
        """Protects the given definition from the GC, if its result is not None."""
        definition.insert_after(cfg_ir.create_conditional_gc_protect(definition, root_node))

    for block in cfg_ir.get_all_blocks(entry_point):
        for definition in block.definitions:
            def_value = cfg_ir.get_def_value(definition)
            if isinstance(def_value, cfg_ir.CreateNode):
                protect_def_from_gc(definition)
            elif (isinstance(def_value, cfg_ir.IndirectFunctionCall)
                  or (isinstance(def_value, cfg_ir.DirectFunctionCall)
                      and (def_value.calling_convention == cfg_ir.JIT_CALLING_CONVENTION
                           or def_value.calling_convention == cfg_ir.JIT_NO_GC_CALLING_CONVENTION
                           or def_value.calling_convention == cfg_ir.MACRO_IO_CALLING_CONVENTION)
                      and def_value.has_value())):
                maybe_protect_def_from_gc(definition)

        if isinstance(block.flow, (cfg_ir.ReturnFlow, cfg_ir.ThrowFlow)):
            block.append_definition(cfg_ir.DeallocateRootNode(root_node))

def elide_gc_protects(entry_point):
    """Tries to elide GC protection values."""
    # We don't need to protect a value from the GC if it is used for the
    # last time _before_ the GC has an opportunity to kick in. To simplify
    # things, we'll do a quick block-based analysis.
    def __may_cause_gc(definition):
        def_value = cfg_ir.get_def_value(definition)
        if isinstance(def_value, cfg_ir.IndirectFunctionCall):
            return True
        elif (isinstance(def_value, cfg_ir.DirectFunctionCall)
              and (def_value.calling_convention == cfg_ir.JIT_CALLING_CONVENTION
                   or def_value.calling_convention == cfg_ir.MACRO_IO_CALLING_CONVENTION)):
            return True
        else:
            return False

    def __get_protected_def(def_or_value):
        value = cfg_ir.get_def_value(def_or_value)
        if cfg_ir.is_call(
                value, target_name=cfg_ir.GC_PROTECT_MACRO_NAME,
                calling_convention=cfg_ir.MACRO_POSITIONAL_CALLING_CONVENTION):
            _, protected_def = value.argument_list[0]
            return protected_def
        elif cfg_ir.is_call(
                value, target_name=cfg_ir.MAYBE_GC_PROTECT_MACRO_NAME,
                calling_convention=cfg_ir.MACRO_POSITIONAL_CALLING_CONVENTION):
            _, protected_def = value.argument_list[1]
            return protected_def
        else:
            return None

    def_blocks = {}
    def __register_def_or_use(definition, block):
        if definition in def_blocks and def_blocks[definition] != block:
            # Definition seems to be used across basic blocks.
            ineligible_defs.add(definition)

        def_blocks[definition] = block

    ineligible_defs = set()
    def_protections = defaultdict(list)
    for block in cfg_ir.get_all_blocks(entry_point):
        no_gc_defs = set()
        block_defs = set()
        first_gc = {}
        last_def_uses = {}
        for i, definition in enumerate(block.definitions):
            if isinstance(definition.value, cfg_ir.Definition):
                # Handling definitions of definitions is complicated and they should already have
                # been expanded at this point. Just mark them as ineligible.
                ineligible_defs.add(definition)
                ineligible_defs.add(definition.value)
                continue

            protected_def = __get_protected_def(definition)
            if protected_def is not None:
                # We just ran into a gc_protect/maybe_gc_protect.
                def_protections[protected_def].append(definition)
                continue

            block_defs.add(definition)
            __register_def_or_use(definition, block)

            for dependency in definition.get_all_dependencies():
                __register_def_or_use(dependency, block)
                last_def_uses[dependency] = i

            if __may_cause_gc(definition):
                for gc_def in no_gc_defs:
                    first_gc[gc_def] = i
                no_gc_defs = set()

            no_gc_defs.add(definition)

        # Mark all branch arguments as ineligible.
        for branch in block.flow.branches():
            ineligible_defs.update(branch.arguments)

        for dependency in block.flow.get_dependencies():
            last_def_uses[dependency] = None

        for definition in block_defs:
            if definition in ineligible_defs:
                # Definition was already ineligible.
                continue

            # Mark `definition` as ineligible if there is a GC definition in the range of
            # definitions (definition, last_def_uses[definition]].
            if definition in first_gc:
                if definition in last_def_uses:
                    last_use = last_def_uses[definition]
                    if last_use is None or first_gc[definition] <= last_use:
                        ineligible_defs.add(definition)

    # Elide all GC protections for definitions which are not in the `ineligible_defs` set.
    for protected, protections in def_protections.items():
        if protected not in ineligible_defs:
            for protect_def in protections:
                protect_def.redefine(cfg_ir.Literal(None))

def optimize(entry_point, jit):
    """Optimizes the control-flow graph defined by the given entry point.
       A potentially altered entry point is returned."""
    optimize_graph_flow(entry_point)
    elide_local_checks(entry_point)
    optimize_graph_flow(entry_point)
    merge_blocks(entry_point)
    eliminate_trivial_phis(entry_point)
    entry_point = cfg_ssa_construction.construct_ssa_form(entry_point)
    if jit.direct_calls_allowed:
        optimize_calls(entry_point, jit)
    cfg_data_structures.optimize_data_structures(entry_point)
    expand_cfg_intrinsics(entry_point, jit)
    yield [("CALL_ARGS", [inline_constants, (entry_point,)])]
    optimize_reads(entry_point)
    simplify_values(entry_point)
    eliminate_unused_definitions(entry_point)
    optimize_graph_flow(entry_point)
    expand_indirect_definitions(entry_point)
    eliminate_unused_definitions(entry_point)
    merge_blocks(entry_point)
    expand_indirect_definitions(entry_point)
    protect_from_gc(entry_point)
    elide_gc_protects(entry_point)
    eliminate_unused_definitions(entry_point)
    raise primitive_functions.PrimitiveFinished(entry_point)