sp-04/src/ch/usi/inf/sp/dom/DominatorTree.java

package ch.usi.inf.sp.dom;

import ch.usi.inf.sp.cfg.BasicBlock;
import ch.usi.inf.sp.graph.DiGraph;

import java.util.HashMap;
import java.util.Map;


public final class DominatorTree extends DiGraph<DominatorTreeNode,DominanceEdge> {

    private Map<BasicBlock,DominatorTreeNode> nodeForBlock;
    private DominatorTreeNode rootNode;


    public DominatorTree() {
        nodeForBlock = new HashMap<>();
    }

    public DominatorTreeNode getNodeForBlock(final BasicBlock block) {
        return nodeForBlock.get(block);
    }

    public DominatorTreeNode setRootBlock(BasicBlock rootBlock) {
        DominatorTreeNode newRootNode = nodeForBlock.get(rootBlock);
        if (newRootNode==null) {
            newRootNode = new DominatorTreeNode(rootBlock);
            nodeForBlock.put(rootBlock, newRootNode);
        }
        if (!getNodes().contains(newRootNode)) {
            addNode(newRootNode);
        }
        rootNode = newRootNode;
        return rootNode;
    }

    public DominatorTreeNode getRoot() {
        return rootNode;
    }

    public DominanceEdge addDominanceEdge(BasicBlock idom, BasicBlock child) {
        DominatorTreeNode idomNode = nodeForBlock.get(idom);
        if (idomNode==null) {
            idomNode = new DominatorTreeNode(idom);
            nodeForBlock.put(idom, idomNode);
        }
        if (!getNodes().contains(idomNode)) {
            addNode(idomNode);
        }
        DominatorTreeNode childNode = nodeForBlock.get(child);
        if (childNode==null) {
            childNode = new DominatorTreeNode(child);
            nodeForBlock.put(child, childNode);
        }
        if (!getNodes().contains(childNode)) {
            addNode(childNode);
        }
        final DominanceEdge edge = new DominanceEdge();
        addEdge(edge);
        connect(idomNode, edge, childNode);
        return edge;
    }

    public String toString() {
        final StringBuffer sb = new StringBuffer("digraph DOM {\n");
        for (final DominatorTreeNode node : getNodes()) {
            if (node==rootNode) {
                sb.append("  " + node + " [style=filled]\n");
            } else {
                sb.append("  " + node + "\n");
            }
        }
        for (final DominanceEdge edge : getEdges()) {
            sb.append("  "+edge+"\n");
        }
        sb.append("}\n");
        return sb.toString();
    }

}