Create classGraph

package es.upv.mist.slicing.graphs;

import com.github.javaparser.ast.CompilationUnit;

import com.github.javaparser.ast.NodeList;

import com.github.javaparser.ast.body.CallableDeclaration;

import com.github.javaparser.ast.body.ConstructorDeclaration;

import com.github.javaparser.ast.body.MethodDeclaration;

import com.github.javaparser.ast.expr.MethodCallExpr;

import com.github.javaparser.ast.expr.ObjectCreationExpr;

import com.github.javaparser.ast.stmt.ExplicitConstructorInvocationStmt;

import com.github.javaparser.ast.visitor.VoidVisitorAdapter;

import com.github.javaparser.resolution.Resolvable;

import com.github.javaparser.resolution.declarations.ResolvedMethodLikeDeclaration;

import es.upv.mist.slicing.graphs.cfg.CFG;

import es.upv.mist.slicing.nodes.GraphNode;

import es.upv.mist.slicing.utils.ASTUtils;

import es.upv.mist.slicing.utils.NodeNotFoundException;

import es.upv.mist.slicing.utils.Utils;

import org.jgrapht.graph.DefaultEdge;

import org.jgrapht.graph.DirectedPseudograph;

import org.jgrapht.nio.dot.DOTExporter;

import java.util.*;

public class ClassGraph extends DirectedPseudograph<CallGraph.Vertex, CallGraph.Edge<?>> implements Buildable<NodeList<CompilationUnit>> {

    private final Map<CallableDeclaration<?>, CFG> cfgMap;

    private final Map<CallableDeclaration<?>, CallGraph.Vertex> vertexDeclarationMap = new IdentityHashMap<>();

    private boolean built = false;

    public ClassGraph(Map<CallableDeclaration<?>, CFG> cfgMap) {

        super(null, null, false);

        this.cfgMap = cfgMap;

    }

    /** Resolve a call to its declaration, by using the call AST nodes stored on the edges. */

    public CallableDeclaration<?> getCallTarget(Resolvable<? extends ResolvedMethodLikeDeclaration> call) {

        return edgeSet().stream()

                .filter(e -> e.getCall() == call)

                .map(this::getEdgeTarget)

                .map(CallGraph.Vertex::getDeclaration)

                .map(CallableDeclaration.class::cast)

                .findFirst()

                .orElse(null);

    }

    @Override

    public void build(NodeList<CompilationUnit> arg) {

        if (isBuilt())

            return;

        buildVertices(arg);

        buildEdges(arg);

        built = true;

    }

    @Override

    public boolean isBuilt() {

        return built;

    }

    /** Find the method and constructor declarations (vertices) in the given list of compilation units. */

    protected void buildVertices(NodeList<CompilationUnit> arg) {

        arg.accept(new VoidVisitorAdapter<Void>() {

            @Override

            public void visit(MethodDeclaration n, Void arg) {

                addDeclaration(n);

                super.visit(n, arg);

            }

            @Override

            public void visit(ConstructorDeclaration n, Void arg) {

                addDeclaration(n);

                super.visit(n, arg);

            }

        }, null);

    }

    protected void addDeclaration(CallableDeclaration<?> n) {

        CallGraph.Vertex v = new CallGraph.Vertex(n);

        vertexDeclarationMap.put(n, v);

        addVertex(v);

    }

    protected boolean addEdge(CallableDeclaration<?> source, CallableDeclaration<?> target, Resolvable<? extends ResolvedMethodLikeDeclaration> call) {

        CallGraph.Edge<?> edge = new CallGraph.Edge<>(call, findGraphNode(call, source));

        return addEdge(vertexDeclarationMap.get(source), vertexDeclarationMap.get(target), edge);

    }

    /** Find the calls to methods and constructors (edges) in the given list of compilation units. */

    protected void buildEdges(NodeList<CompilationUnit> arg) {

        arg.accept(new VoidVisitorAdapter<Void>() {

            private final Deque<CallableDeclaration<?>> declStack = new LinkedList<>();

            // ============ Method declarations ===========

            // There are some locations not considered, which may lead to an error in the stack.

            // 1. Method calls in non-static field initializations are assigned to all constructors of that class

            // 2. Method calls in static field initializations are assigned to the static block of that class

            @Override

            public void visit(MethodDeclaration n, Void arg) {

                declStack.push(n);

                super.visit(n, arg);

                declStack.pop();

            }

            @Override

            public void visit(ConstructorDeclaration n, Void arg) {

                declStack.push(n);

                super.visit(n, arg);

                declStack.pop();

            }

            // =============== Method calls ===============

            @Override

            public void visit(MethodCallExpr n, Void arg) {

                n.resolve().toAst().ifPresent(decl -> addEdge(declStack.peek(), decl, n));

                super.visit(n, arg);

            }

            @Override

            public void visit(ObjectCreationExpr n, Void arg) {

                n.resolve().toAst().ifPresent(decl -> addEdge(declStack.peek(), decl, n));

                super.visit(n, arg);

            }

            @Override

            public void visit(ExplicitConstructorInvocationStmt n, Void arg) {

                n.resolve().toAst().ifPresent(decl -> addEdge(declStack.peek(), decl, n));

                super.visit(n, arg);

            }

        }, null);

    }

    /** Locates the node in the collection of CFGs that contains the given call. */

    protected GraphNode<?> findGraphNode(Resolvable<? extends ResolvedMethodLikeDeclaration> n, CallableDeclaration<?> declaration) {

        for (GraphNode<?> node : cfgMap.get(declaration).vertexSet())

            if (node.containsCall(n))

                return node;

        throw new NodeNotFoundException("call " + n + " could not be located!");

    }

    /** Creates a graph-appropriate DOT exporter. */

    public DOTExporter<CallableDeclaration<?>, CallGraph.Edge<?>> getDOTExporter() {

        DOTExporter<CallableDeclaration<?>, CallGraph.Edge<?>> dot = new DOTExporter<>();

        dot.setVertexAttributeProvider(decl -> Utils.dotLabel(decl.getDeclarationAsString(false, false, false)));

        dot.setEdgeAttributeProvider(edge -> Utils.dotLabel(edge.getCall().toString()));

        return dot;

    }

    /** A vertex containing the declaration it represents. It only exists because

     *  JGraphT relies heavily on equals comparison, which may not be correct in declarations. */

    public static class Vertex {

        protected final CallableDeclaration<?> declaration;

        public Vertex(CallableDeclaration<?> declaration) {

            assert declaration instanceof ConstructorDeclaration || declaration instanceof MethodDeclaration;

            this.declaration = declaration;

        }

        /** The declaration represented by this node. */

        public CallableDeclaration<?> getDeclaration() {

            return declaration;

        }

        @Override

        public int hashCode() {

            return Objects.hash(declaration, declaration.getRange());

        }

        @Override

        public boolean equals(Object obj) {

            return obj instanceof CallGraph.Vertex && ASTUtils.equalsWithRangeInCU(((CallGraph.Vertex) obj).declaration, declaration);

        }

        @Override

        public String toString() {

            return super.toString();

        }

    }

    /** An edge containing the call it represents, and the graph node that contains it. */

    public static class Edge<T extends Resolvable<? extends ResolvedMethodLikeDeclaration>> extends DefaultEdge {

        protected final T call;

        protected final GraphNode<?> graphNode;

        public Edge(T call, GraphNode<?> graphNode) {

            assert call instanceof MethodCallExpr || call instanceof ObjectCreationExpr || call instanceof ExplicitConstructorInvocationStmt;

            this.call = call;

            this.graphNode = graphNode;

        }

        /** The call represented by this edge. Using the {@link CallGraph} it can be effortlessly resolved. */

        public T getCall() {

            return call;

        }

        /** The graph node that contains the call represented by this edge. */

        public GraphNode<?> getGraphNode() {

            return graphNode;

        }

        @Override

        public String toString() {

            return String.format("%s -%d-> %s",

                    ((CallableDeclaration<?>) getSource()).getDeclarationAsString(false, false, false),

                    graphNode.getId(),

                    ((CallableDeclaration<?>) getTarget()).getDeclarationAsString(false, false, false));

        }

    }

}