sp-05/src/ch/usi/inf/sp/callgraph/renderer/DotGraph.java

package ch.usi.inf.sp.callgraph.renderer;

import ch.usi.inf.sp.callgraph.*;

import java.util.*;
import java.util.stream.Collectors;

public class DotGraph {
    // use map[e->e] instead of set[e] to allow fetching equal element
    private final DuplicateAwareSet<DotClassNode> classNodes = new DuplicateAwareSet<>();
    private final DuplicateAwareSet<DotMethodNode> methodNodes = new DuplicateAwareSet<>();
    private final List<DotEdge> otherEdges = new ArrayList<>();
    private final Set<DotEdge> classToMethodEdges = new HashSet<>();

    private static String formatProperties(Map<String, String> properties) {
        return properties.entrySet().stream()
                .map((e) -> String.format("%s=\"%s\"", e.getKey(), e.getValue()))
                .collect(Collectors.joining(","));
    }

    private static <E> Map<E, Integer> enumerate(Set<E> set) {
        final Map<E, Integer> map = new HashMap<>();
        int index = 0;

        for (final E element : set) {
            map.put(element, index++);
        }

        return Collections.unmodifiableMap(map);
    }

    private DotClassNode considerClass(final ClassType classType) {
        return classNodes.addIfNew(new DotClassNode(classType), (addedClassNode) -> {
            if (addedClassNode.isResolved()) {
                for (final Method m : classType.getMethods()) {
                    methodNodes.addIfNew(
                            new DotMethodNode(new MethodAdapter(m)),
                            (addedMethodNode) -> classToMethodEdges.add(new ClassToMethodEdge(addedClassNode, addedMethodNode))
                    );
                }

                for (final ClassType c : classType.getInterfaces()) {
                    final DotClassNode interfaceNode = considerClass(c);
                    otherEdges.add(new TypeHierarchyEdge(addedClassNode, interfaceNode));
                }
            }
        });
    }

    private void scanMethods() {

        for (final DotMethodNode fromMethodNode : methodNodes.getAll()) {
            final Method m = fromMethodNode.getMethodLike().getMethod();
            if (m == null) continue; // we care about actual methods only, not callsites

            for (final CallSite c : m.getCallSites()) {
                final DotMethodNode toMethodNode = methodNodes.addIfNew(new DotMethodNode(new CallSiteAdapter(c)));
                otherEdges.add(new InvokeEdge(fromMethodNode, toMethodNode, c.getOpcode(), true));

                for (final ClassType p : c.getPossibleTargetClasses()) {
                    final DotMethodNode toPossibleMethodNode = methodNodes.addIfNew(new DotMethodNode(new PossibleTargetAdapter(c, p)));
                    otherEdges.add(new InvokeEdge(fromMethodNode, toPossibleMethodNode, c.getOpcode(), false));

                    // if the target is a new node, the class of that node might not have been analyzed,
                    // thus class to method edges might not have been added. We re-add the class to method edge, and
                    // duplicates are handled by the Set collection.
                    final DotClassNode toPossibleClassNode = classNodes.addIfNew(new DotClassNode(p));
                    classToMethodEdges.add(new ClassToMethodEdge(toPossibleClassNode, toPossibleMethodNode));
                }
            }
        }
    }

    public void build(final ClassHierarchy hierarchy) {
        for (final Type type : hierarchy.getTypes()) {
            if (type instanceof ClassType) {
                final ClassType classType = (ClassType) type;
                final DotClassNode node = considerClass(classType);

                final ClassType superClassType = classType.getSuperClass();
                if (superClassType != null) {
                    final DotClassNode superNode = considerClass(superClassType);
                    otherEdges.add(new TypeHierarchyEdge(node, superNode));
                }
            }
        }

        scanMethods();
    }

    public String toDot() {
        final StringBuilder s = new StringBuilder();
        s.append("digraph CallGraph {\n");
        s.append("rankdir=\"BT\"\n");

        final Set<DotNode> nodes = new HashSet<>();
        nodes.addAll(classNodes.getAll());
        nodes.addAll(methodNodes.getAll());

        final Map<DotNode, Integer> nodeToId = enumerate(nodes);

        for (final Map.Entry<DotNode, Integer> e : nodeToId.entrySet()) {
            s.append(buildNode(e.getKey(), e.getValue()));
        }

        s.append('\n');

        final List<DotEdge> edges = new ArrayList<>(otherEdges);
        edges.addAll(classToMethodEdges);

        for (final DotEdge e : edges) {
            s.append(buildEdge(e, nodeToId.get(e.getFrom()), nodeToId.get(e.getTo())));
        }

        s.append("}\n");
        return s.toString();
    }

    private String buildNode(DotNode node, int nodeId) {
        return String.format("node%d [%s];\n", nodeId, formatProperties(node.getProperties()));
    }

    private String buildEdge(DotEdge edge, int idFrom, int idTo) {
        return String.format("node%d %s node%d [%s];\n", idFrom, edge.getConnector(), idTo, formatProperties(edge.getProperties()));
    }
}