"""Computes dominator trees for control-flow graphs."""

from collections import defaultdict

def sort_postorder(entry_point):
    """Produces a postorder traversal of the graph with the given block as entry point."""
    processed = set()
    results = []
    def __sort_postorder_step(block):
        if block in processed:
            return

        processed.add(block)
        for branch in block.flow.branches():
            __sort_postorder_step(branch.block)

        results.append(block)

    __sort_postorder_step(entry_point)
    return results

def get_all_predecessor_blocks(entry_point):
    """Creates a mapping of blocks to their direct predecessors for every block in the control-flow
       graph defined by the given entry point."""
    results = defaultdict(set)
    processed = set()
    def __find_predecessors_step(block):
        if block in processed:
            return

        processed.add(block)
        for branch in block.flow.branches():
            target_block = branch.block
            results[target_block].add(block)
            __find_predecessors_step(target_block)

    __find_predecessors_step(entry_point)
    return results

def intersect_immediate_dominators(block1, block2, idoms, postorder_nums):
    """Computes the intersection of the given blocks' immediate dominators."""
    finger1 = block1
    finger2 = block2
    while finger1 != finger2:
        while postorder_nums[finger1] < postorder_nums[finger2]:
            finger1 = idoms[finger1]

        while postorder_nums[finger2] < postorder_nums[finger1]:
            finger2 = idoms[finger2]
    return finger1

def get_immediate_dominators(entry_point):
    """Computes the immediate dominators of the control-flow graph defined by the given entry
       point."""
    # Based on "A Simple, Fast Dominance Algorithm" by
    # Keith D. Cooper, Timothy J. Harvey, and Ken Kennedy
    # (http://www.cs.rice.edu/~keith/Embed/dom.pdf)

    predecessor_map = get_all_predecessor_blocks(entry_point)
    idoms = {}
    postorder_sort = sort_postorder(entry_point)
    postorder_nums = {}
    for i, elem in enumerate(postorder_sort):
        postorder_nums[elem] = i
        idoms[elem] = None

    idoms[entry_point] = entry_point

    changed = True
    while changed:
        changed = False
        for block in reversed(postorder_sort):
            if block == entry_point:
                continue

            new_idom = None
            for pred in predecessor_map[entry_point]:
                if pred not in postorder_nums:
                    continue

                if new_idom is None:
                    new_idom = pred
                elif idoms[pred] is not None:
                    new_idom = intersect_immediate_dominators(pred, new_idom, idoms, postorder_nums)

            if idoms[block] is not new_idom:
                idoms[block] = new_idom
                changed = True

    return idoms

def get_dominator_tree(entry_point):
    """Constructs the dominator tree for the control-flow graph defined by the given entry point."""
    return DominatorTree(get_immediate_dominators(entry_point))

class DominatorTree(object):
    """A data structure that represents a dominator tree."""
    def __init__(self, immediate_dominators):
        self.immediate_dominators = immediate_dominators
        self.dominator_sets = {}

    def get_immediate_dominator(self, block):
        """Gets the given block's immediate dominator."""
        return self.immediate_dominators[block]

    def get_dominators(self, block):
        """Gets the given block's set of all dominators."""
        if block in self.dominator_sets:
            return self.dominator_sets
        else:
            idom = self.get_immediate_dominator(block)
            if idom == block:
                results = set([block])
            else:
                results = set(self.get_dominators(idom))
                results.add(block)

            self.dominator_sets[block] = results
            return results

    def dominates(self, dominator, dominated):
        """Tests if the first block dominates the second."""
        return dominator in self.get_dominators(dominated)