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.BooleanLiteralExpr;
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;

/**
 * Visits a compilation unit and extracts the control flow graph, registering all the data in the argument to all
 * functions: {@link CFG}. This class contains no information, all data is saved in the control flow graph.
 */
public class Visitador extends VoidVisitorAdapter<CFG> {
	@Override
	public void visit(MethodDeclaration methodDeclaration, CFG graph) {
		graph.beginMethod();
		super.visit(methodDeclaration, graph);
		graph.finishMethod();
	}

	/**
	 * Return statements are connected to the previous statement and to the end node of the graph.
	 * Schema: {@code -> return -> Stop}
	 */
	@Override
	public void visit(ReturnStmt n, CFG arg) {
		arg.connectTo(n);
		arg.connectTo(arg.getEndNode());
		arg.clearNode();
	}

	/** If statements are set up as a block, with a condition and supports any configuration, including missing else,
	 * empty then/else blocks. */
	@Override
	public void visit(IfStmt n, CFG graph) {
		// prev --> cond
		Node ifStart = n.getCondition();
		graph.beginBlock(n, ifStart);
		graph.connectTo(ifStart);
		// cond --> then (if true)
		graph.setNextLabel(true);
		n.getThenStmt().accept(this, graph);
		// then --> next (saved for later)
		List<Node> newPrev = new LinkedList<>(graph.getNodeList());
		// Else processing (either connect cond --> else --> next or cond --> next)
		graph.setNode(ifStart);
		graph.setNextLabel(false);
		if (n.getElseStmt().isPresent()) {
			// cond --> else (if false)
			n.getElseStmt().get().accept(this, graph);
			// else --> next
			newPrev.addAll(graph.getNodeList());
		} else {
			// cond --> next (if false)
			newPrev.add(ifStart);
		}
		// Set ends of then/else as the new list
		graph.setNodeList(newPrev);
		graph.endBlock();
	}

	/** Registers the label used for the loop statement in the break and continue sections of the CFG */
	@Override
	public void visit(LabeledStmt n, CFG graph) {
		graph.beginBreakSection(n.getLabel());
		graph.beginContinueSection(n.getLabel());
		super.visit(n, graph);
		graph.endContinueSection(n.getLabel());
		graph.endBreakSection(n.getLabel());
	}

	/** While statements are set up as a block. They also generate an unlabeled break and continue section.
	 * The schema followed is {@code -> cond -> body -> cond ->} */
	@Override
	public void visit(WhileStmt n, CFG graph) {
		Node whileStart = n.getCondition();
		// Begin block/sections
		graph.beginBlock(n, whileStart);
		graph.beginBreakSection();
		graph.beginContinueSection(whileStart);
		// prev --> cond
		graph.connectTo(whileStart);
		// cond --> body (if true)
		graph.setNextLabel(true);
		n.getBody().accept(this, graph);
		// body --> cond
		graph.connectTo(whileStart);
		// cond --> next (if false)
		graph.setNextLabel(false);
		// End block/sections
		graph.endContinueSection();
		graph.endBreakSection();
		graph.endBlock();
	}

	/** Do while statement, set up as block. Generates break and continue section. Schema:
	 * {@code -> body -> cond -> body -> cond ->} */
	@Override
	public void visit(DoStmt n, CFG graph) {
		Node condition = n.getCondition();
		// Begin block/section
		graph.beginBlock(n, condition);
		graph.beginBreakSection();
		graph.beginContinueSection(condition);
		// prev --> body
		// cond --> body (if true)
		graph.appendNode(condition);
		n.getBody().accept(this, graph);
		// body --> cond
		graph.connectTo(condition);
		// cond --> next (if false)
		graph.setNextLabel(false);
		// End block/section
		graph.endContinueSection();
		graph.endBreakSection();
		graph.endBlock();
	}

	/** For statement, set up as a block. Generates break and continue section. If the condition is empty, it will be
	 * set as a {@code true} literal expression. Schema: {@code -> initializer(s) -> cond -> body -> update(s) -> cond -> } */
	@Override
	public void visit(ForStmt n, CFG graph) {
		// Begin block and break sections (continue section needs the comparison, see below)
		graph.beginBlock(n, n.getCompare().orElse(null));
		graph.beginBreakSection();
		// prev --> initializer(s), chained in serial
		for (Expression e : n.getInitialization())
			graph.connectTo(e);
		// Set the comparison expression to "true" if the expression was empty, as a safeguard and to ease the treatment
		// of cases where the loop is of the form "for (;;) ;", which would require special processing.
		assert n.getTokenRange().isPresent();
		if (!n.getCompare().isPresent())
			n.setCompare(new BooleanLiteralExpr(n.getTokenRange().get(), true));
		// Begin continue section (with the comparison on hand)
		graph.beginContinueSection(n.getCompare().get());
		// initializer(s) --> cond
		graph.connectTo(n.getCompare().get());
		// cond --> body (if true)
		graph.setNextLabel(true);
		n.getBody().accept(this, graph);
		// body --> update(s), chained in serial
		for (Expression e : n.getUpdate())
			graph.connectTo(e);
		// update(s) --> comp
		graph.connectTo(n.getCompare().get());
		// comp --> next (if false)
		graph.setNextLabel(false);
		// End block/section (all of them)
		graph.endContinueSection();
		graph.endBreakSection();
		graph.endBlock();
	}

	/** Foreach statement, set up as a block. Generates break and continue sections. Schema: {@code -> cond -> body -> cond ->} */
	@Override
	public void visit(ForeachStmt n, CFG graph) {
		// Create a copy of the foreach statement to use as the condition node. The alternative would be to "emulate" the
		// foreach statement with a for statement.
		ForeachStmt copy = new ForeachStmt(n.getTokenRange().orElse(null), n.getVariable(), n.getIterable(), new EmptyStmt());
		// Begin block/section
		graph.beginBlock(n, copy);
		graph.beginBreakSection();
		graph.beginContinueSection(copy);
		// prev --> cond
		graph.connectTo(copy);
		// cond --> body (if true)
		graph.setNextLabel(true);
		n.getBody().accept(this, graph);
		// body --> cond
		graph.connectTo(copy);
		// cond --> next (if false)
		graph.setNextLabel(false);
		graph.endContinueSection();
		graph.endBreakSection();
		graph.endBlock();
	}

	/** Switch entry, considered part of the condition of the switch. */
	@Override
	public void visit(SwitchEntryStmt n, CFG graph) {
		// Case header (prev -> case EXPR)
		Node caseNode = new SwitchEntryStmt(n.getTokenRange().orElse(null), n.getLabel().orElse(null), new NodeList<>());
		graph.connectTo(caseNode);
		// Case body (case EXPR --> body)
		n.getStatements().accept(this, graph);
		// body --> next
	}

	/** Switch statement, set up as a block. Generates an unlabeled break section. The statement begins with an expression,
	 * which is compared to each of the cases, entering through the first one that matches. Each case may end in a break
	 * or return statement, or not. When it doesn't the control falls through the next cases. If the switch doesn't have
	 * a default entry, the */
	@Override
	public void visit(SwitchStmt n, CFG graph) {
		// Link previous statement to the switch's selector
		Node selectorNode = n.getSelector();
		// Begin block and break section
		graph.beginBlock(n, selectorNode);
		graph.beginBreakSection();
		// prev --> expr
		graph.connectTo(selectorNode);
		// expr --> each case (fallthrough by default, so case --> case too)
		for (SwitchEntryStmt entry : n.getEntries()) {
			entry.accept(this, graph); // expr && prev case --> case --> next case
			graph.appendNode(selectorNode); // expr --> next case
		}
		// The next statement will be linked to:
		//		1. All break statements that broke from the switch (done with break section)
		// 		2. If the switch doesn't have a default statement, the switch's selector (already present)
		// 		3. If the last entry doesn't break, to the last statement (present already)
		// If the last case is a default case, remove the selector node from the list of nodes (see 2)
		if (!n.getEntries().get(n.getEntries().size() - 1).getLabel().isPresent())
			graph.removeNode(selectorNode);
		// End block and break section
		graph.endBreakSection();
		graph.endBlock();
	}

	@Override
	public void visit(BreakStmt n, CFG graph) {
		graph.connectBreak(n); // prev --> break --> end of loop
	}

	@Override
	public void visit(ContinueStmt n, CFG graph) {
		graph.connectContinue(n); // prev --> continue --> loop condition
	}

	@Override
	public void visit(ExpressionStmt es, CFG graph) {
		graph.connectTo(es); // prev --> es --> next
	}
}