CFGVisitor.java

package tfm.visitors;

import com.github.javaparser.ast.Node;
import com.github.javaparser.ast.NodeList;
import com.github.javaparser.ast.body.MethodDeclaration;
import com.github.javaparser.ast.body.VariableDeclarator;
import com.github.javaparser.ast.expr.*;
import com.github.javaparser.ast.stmt.*;
import com.github.javaparser.ast.visitor.GenericVisitor;
import com.github.javaparser.ast.visitor.VoidVisitor;
import com.github.javaparser.ast.visitor.VoidVisitorAdapter;
import sun.rmi.runtime.Log;
import tfm.graphs.CFGGraph;
import tfm.nodes.CFGNode;
import tfm.utils.Logger;
import tfm.utils.Utils;

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

public class CFGVisitor extends VoidVisitorAdapter<Void> {

    private CFGGraph graph;

    private Queue<CFGNode> lastParentNodes;

    public CFGVisitor(CFGGraph graph) {
        this.graph = graph;
        this.lastParentNodes = Collections.asLifoQueue(
                new ArrayDeque<>(
                        Collections.singletonList(graph.getRootNode())
                )
        );
    }

    @Override
    public void visit(ExpressionStmt expressionStmt, Void arg) {
        String expression = expressionStmt.toString().replace("\"", "\\\"");

        CFGNode nextNode = addNodeAndArcs(expression, expressionStmt);

        lastParentNodes.add(nextNode);
    }

//    @Override
//    public void visit(VariableDeclarationExpr variableDeclarationExpr, Void arg) {
//        CFGNode<String> nextNode = addNodeAndArcs(variableDeclarationExpr.toString());
//
//        lastParentNodes.add(nextNode);
//
//        Logger.log(variableDeclarationExpr);
//
//        super.visit(variableDeclarationExpr, arg);
//    }

    @Override
    public void visit(IfStmt ifStmt, Void arg) {
        CFGNode ifCondition = addNodeAndArcs(
                String.format("if (%s)", ifStmt.getCondition().toString()),
                ifStmt
        );

        lastParentNodes.add(ifCondition);

        // Visit "then"
        ifStmt.getThenStmt().accept(this, arg);

        Queue<CFGNode> lastThenNodes = new ArrayDeque<>(lastParentNodes);

        if (ifStmt.hasElseBranch()) {
            lastParentNodes.clear();
            lastParentNodes.add(ifCondition); // Set if nodes as root

            ifStmt.getElseStmt().get().accept(this, arg);

            lastParentNodes.addAll(lastThenNodes);
        } else {
            lastParentNodes.add(ifCondition);
        }
    }

    @Override
    public void visit(WhileStmt whileStmt, Void arg) {
        CFGNode whileCondition = addNodeAndArcs(
                String.format("while (%s)", whileStmt.getCondition().toString()),
                whileStmt
        );

        lastParentNodes.add(whileCondition);

        whileStmt.getBody().accept(this, arg);

        while (!lastParentNodes.isEmpty()) {
            graph.addControlFlowEdge(lastParentNodes.poll(), whileCondition);
        }

        lastParentNodes.add(whileCondition);
    }

    @Override
    public void visit(DoStmt doStmt, Void arg) {
        BlockStmt body = Utils.blockWrapper(doStmt.getBody());

        body.accept(this, arg);

        CFGNode doWhileNode = addNodeAndArcs(
                String.format("while (%s)", doStmt.getCondition()),
                doStmt
        );

        if (!body.isEmpty()) {
            Statement firstBodyStatement = body.getStatement(0);

            graph.findNodeByStatement(firstBodyStatement)
                    .ifPresent(node -> graph.addControlFlowEdge(doWhileNode, node));
        }

        lastParentNodes.add(doWhileNode);
    }

    @Override
    public void visit(ForStmt forStmt, Void arg) {
//        String inizialization = forStmt.getInitialization().stream()
//                .map(Node::toString)
//                .collect(Collectors.joining(","));
//
//        String update = forStmt.getUpdate().stream()
//                .map(Node::toString)
//                .collect(Collectors.joining(","));

        Expression comparison = forStmt.getCompare().orElse(new BooleanLiteralExpr(true));

        forStmt.getInitialization().forEach(expression -> new ExpressionStmt(expression).accept(this, null));

        CFGNode forNode = addNodeAndArcs(
                String.format("for (;%s;)", comparison),
                forStmt
        );

        lastParentNodes.add(forNode);

        BlockStmt body = Utils.blockWrapper(forStmt.getBody());

        forStmt.getUpdate().forEach(body::addStatement);

        body.accept(this, arg);

        while (!lastParentNodes.isEmpty()) {
            graph.addControlFlowEdge(lastParentNodes.poll(), forNode);
        }

        lastParentNodes.add(forNode);
    }

    @Override
    public void visit(ForEachStmt forEachStmt, Void arg) {
        CFGNode foreachNode = addNodeAndArcs(
                String.format("for (%s : %s)", forEachStmt.getVariable(), forEachStmt.getIterable()),
                forEachStmt
        );

        lastParentNodes.add(foreachNode);

        forEachStmt.getBody().accept(this, arg);

        while (!lastParentNodes.isEmpty()) {
            graph.addControlFlowEdge(lastParentNodes.poll(), foreachNode);
        }

        lastParentNodes.add(foreachNode);
    }

    @Override
    public void visit(SwitchStmt switchStmt, Void arg) {
        CFGNode switchNode = addNodeAndArcs(
                String.format("switch (%s)", switchStmt.getSelector()),
                switchStmt
        );

        lastParentNodes.add(switchNode);

        List<CFGNode> lastEntryParents = new ArrayList<>();

        switchStmt.getEntries().forEach(entry -> {
            Optional<BreakStmt> entryBreak = entry.findFirst(BreakStmt.class, breakStmt -> {
                Optional<Node> parent = breakStmt.getParentNode();

                return parent.isPresent() && parent.get()   .equals(entry);
            });

            new BlockStmt(entry.getStatements()).accept(this, arg);

            if (entryBreak.isPresent()) {
                while (!lastParentNodes.isEmpty()) {
                    lastEntryParents.add(lastParentNodes.poll());
                }
            }

            lastParentNodes.add(switchNode);
        });

        lastParentNodes.clear();
        lastParentNodes.addAll(lastEntryParents);
    }

    @Override
    public void visit(BreakStmt breakStmt, Void arg) {

    }

    @Override
    public void visit(ContinueStmt continueStmt, Void arg) {

    }

    @Override
    public void visit(MethodDeclaration methodDeclaration, Void arg) {
        super.visit(methodDeclaration, arg);

        addNodeAndArcs("Stop", new EmptyStmt());
    }

    private CFGNode addNodeAndArcs(String nodeData, Statement statement) {
        CFGNode node = graph.addNode(nodeData, statement);

        CFGNode parent = lastParentNodes.poll(); // ALWAYS exists a parent
        graph.addControlFlowEdge(parent, node);

        while (!lastParentNodes.isEmpty()) {
            parent = lastParentNodes.poll();
            graph.addControlFlowEdge(parent, node);
        }

        return node;
    }
}