ltd-graph-builder/src/main/java/grafos/Visitador.java

package grafos;

import com.github.javaparser.ast.Node;
import com.github.javaparser.ast.NodeList;
import com.github.javaparser.ast.body.MethodDeclaration;
import com.github.javaparser.ast.expr.Expression;
import com.github.javaparser.ast.stmt.*;
import com.github.javaparser.ast.visitor.VoidVisitorAdapter;

import java.util.LinkedList;
import java.util.List;
import java.util.Stack;

public class Visitador extends VoidVisitorAdapter<CFG> {

	private Stack<SwitchStmt> switchStack = new Stack<>();
	private List<BreakStmt> brokenStmts = new LinkedList<>();

	@Override
	public void visit(MethodDeclaration methodDeclaration, CFG graph) {
		graph.beginMethod();
		super.visit(methodDeclaration, graph);
		graph.finishMethod();
	}

	@Override
	public void visit(ReturnStmt n, CFG arg) {
		arg.connectTo(n);
		arg.connectTo(arg.getEndNode());
		arg.clearNode();
	}

	@Override
	public void visit(IfStmt n, CFG graph) {
		// Connect to the condition
		Node ifStart = n.getCondition();
		graph.beginBlock(n, ifStart);
		graph.connectTo(ifStart);
		// Connect condition to then block
		graph.setNextLabel(true);
		n.getThenStmt().accept(this, graph);
		// Save end of then block
		List<Node> newPrev = new LinkedList<>(graph.getNodeList());
		// Connect condition to else block
		graph.setNode(ifStart);
		graph.setNextLabel(false);
		if (n.getElseStmt().isPresent()) {
			n.getElseStmt().get().accept(this, graph);
			newPrev.addAll(graph.getNodeList());
		} else {
			newPrev.add(ifStart);
		}
		// Set ends of then/else as the new list
		graph.setNodeList(newPrev);
		graph.endBlock();
	}

	@Override
	public void visit(WhileStmt n, CFG graph) {
		Node whileStart = n.getCondition();
		graph.beginBlock(n, whileStart);
		graph.connectTo(whileStart);
		graph.setNextLabel(true);
		n.getBody().accept(this, graph);
		graph.connectTo(whileStart);
		graph.setNextLabel(false);
		graph.endBlock();
	}

	@Override
	public void visit(DoStmt n, CFG graph) {
		Node condition = n.getCondition();
		graph.beginBlock(n, condition);
		graph.appendNode(condition);
		n.getBody().accept(this, graph);
		graph.connectTo(condition);
		graph.setNextLabel(false);
		graph.endBlock();
	}

	@Override
	public void visit(ForStmt n, CFG graph) {
		graph.beginBlock(n, n.getCompare().orElse(null));
		// Initialization expressions
		if (n.getInitialization() != null) {
			for (Expression e : n.getInitialization()) {
				graph.connectTo(e);
			}
		}
		// Comparison expression
		// TODO: shortcut conditions (||, &&)
		if (n.getCompare().isPresent()) {
			graph.connectTo(n.getCompare().get());
			graph.setNextLabel(true);
		}
		// Loop body
		n.getBody().accept(this, graph);
		// Update expressions
		if (n.getUpdate() != null)
			for (Expression e : n.getUpdate())
				graph.connectTo(e);
		// Connect to comparison expression
		// Set comparison as last possible statement
		if (n.getCompare().isPresent()) {
			graph.connectTo(n.getCompare().get());
			graph.setNextLabel(false);
		} else {
			// TODO: connect last update / end of loop with first element of loop body
			// There is no comparison, can't exit the loop TODO implement break and continue as way to exit
			graph.clearNode();
		}
		graph.endBlock();
	}

	@Override
	public void visit(ForeachStmt n, CFG graph) {
		ForeachStmt copy = new ForeachStmt(n.getTokenRange().orElse(null), n.getVariable(), n.getIterable(), new EmptyStmt());
		graph.beginBlock(n, copy);
		graph.connectTo(copy);
		graph.setNextLabel(true);
		n.getBody().accept(this, graph);
		graph.connectTo(copy);
		graph.setNextLabel(false);
		graph.endBlock();
	}

	@Override
	public void visit(SwitchEntryStmt n, CFG graph) {
		// Case header
		Node caseNode = new SwitchEntryStmt(n.getTokenRange().orElse(null), n.getLabel().orElse(null), new NodeList<>());
		graph.connectTo(caseNode);
		// Case body
		n.getStatements().accept(this, graph);
	}

	@Override
	public void visit(SwitchStmt n, CFG graph) {
		// Link previous statement to the switch's selector
		Node selectorNode = n.getSelector();
		graph.beginBlock(n, selectorNode);
		graph.connectTo(selectorNode);
		// Analyze switch's cases
		switchStack.push(n);
		for (SwitchEntryStmt entry : n.getEntries()) {
			entry.accept(this, graph);
			graph.appendNode(selectorNode);
		}
		switchStack.pop();
		// The next statement will be linked to:
		//		1. All break statements that broke from the switch
		// 		2. If the switch doesn't have a default statement, the switch's selector
		// 		3. If the last entry doesn't break, to the last statement
		graph.removeNode(selectorNode);
		graph.appendNodes(brokenStmts); // 1
		brokenStmts.clear();
		if (n.getEntries().get(n.getEntries().size() - 1).getLabel().isPresent())
			graph.appendNode(selectorNode); // 2
		graph.endBlock();
	}

	@Override
	public void visit(BreakStmt n, CFG graph) {
		if (!switchStack.isEmpty() && !n.getLabel().isPresent()) {
			graph.connectTo(n, false);
			brokenStmts.add(n);
			graph.clearNode();
			return;
		}
		super.visit(n, graph);
	}

	@Override
	public void visit(ExpressionStmt es, CFG graph) {
		graph.connectTo(es);
	}
}