GeneralWorker.java

package edu.udel.cis.vsl.civl.transform.common;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;

import edu.udel.cis.vsl.abc.ast.IF.AST;
import edu.udel.cis.vsl.abc.ast.IF.ASTFactory;
import edu.udel.cis.vsl.abc.ast.entity.IF.Entity;
import edu.udel.cis.vsl.abc.ast.entity.IF.Function;
import edu.udel.cis.vsl.abc.ast.entity.IF.OrdinaryEntity;
import edu.udel.cis.vsl.abc.ast.node.IF.ASTNode;
import edu.udel.cis.vsl.abc.ast.node.IF.SequenceNode;
import edu.udel.cis.vsl.abc.ast.node.IF.declaration.DeclarationNode;
import edu.udel.cis.vsl.abc.ast.node.IF.declaration.FunctionDefinitionNode;
import edu.udel.cis.vsl.abc.ast.node.IF.declaration.VariableDeclarationNode;
import edu.udel.cis.vsl.abc.ast.node.IF.expression.CastNode;
import edu.udel.cis.vsl.abc.ast.node.IF.expression.ExpressionNode;
import edu.udel.cis.vsl.abc.ast.node.IF.expression.ExpressionNode.ExpressionKind;
import edu.udel.cis.vsl.abc.ast.node.IF.expression.FunctionCallNode;
import edu.udel.cis.vsl.abc.ast.node.IF.expression.IdentifierExpressionNode;
import edu.udel.cis.vsl.abc.ast.node.IF.expression.OperatorNode;
import edu.udel.cis.vsl.abc.ast.node.IF.expression.OperatorNode.Operator;
import edu.udel.cis.vsl.abc.ast.node.IF.statement.BlockItemNode;
import edu.udel.cis.vsl.abc.ast.node.IF.statement.CompoundStatementNode;
import edu.udel.cis.vsl.abc.ast.node.IF.statement.ExpressionStatementNode;
import edu.udel.cis.vsl.abc.ast.node.IF.statement.ForLoopInitializerNode;
import edu.udel.cis.vsl.abc.ast.node.IF.statement.LoopNode;
import edu.udel.cis.vsl.abc.ast.node.IF.statement.StatementNode;
import edu.udel.cis.vsl.abc.ast.node.IF.type.FunctionTypeNode;
import edu.udel.cis.vsl.abc.ast.node.IF.type.TypeNode;
import edu.udel.cis.vsl.abc.ast.node.IF.type.TypeNode.TypeNodeKind;
import edu.udel.cis.vsl.abc.ast.type.IF.StandardBasicType.BasicTypeKind;
import edu.udel.cis.vsl.abc.ast.type.IF.Type;
import edu.udel.cis.vsl.abc.ast.type.IF.Type.TypeKind;
import edu.udel.cis.vsl.abc.token.IF.Source;
import edu.udel.cis.vsl.abc.token.IF.SyntaxException;
import edu.udel.cis.vsl.abc.transform.IF.NameTransformer;
import edu.udel.cis.vsl.abc.transform.IF.Transform;
import edu.udel.cis.vsl.civl.model.IF.CIVLSyntaxException;

/**
 * The general transformer performs the following transformations:
 * 
 * <ul>
 * <li>malloc(...) to $malloc($gen_root, ...); where $gen_root is the root scope
 * of this AST (which is reserved as a relative root if other transformers make
 * this AST as part of another)</li>
 * <li>arguments of the main function argc and argv become input variables</li>
 * <li>static variables are all moved to the root scope</li>
 * </ul>
 * 
 * @author zmanchun
 *
 */
public class GeneralWorker extends BaseWorker {

	private final static String MALLOC = "malloc";
	final static String GENERAL_ROOT = "$gen_root";
	private final static String MAX_ARGC = "10";
	private final static String separator = "$";

	private final static String INPUT_PREFIX = "CIVL_";
	private int static_var_count = 0;

	private String argvName;
	private String newArgvName;
	private StatementNode argcAssumption = null;
	private Source mainSource;
	/**
	 * static variable declaration nodes of this AST
	 */
	private List<VariableDeclarationNode> static_variables = new LinkedList<>();

	public GeneralWorker(ASTFactory astFactory) {
		super("GeneralTransformer", astFactory);
		this.identifierPrefix = "$gen_";
	}

	@Override
	public AST transform(AST unit) throws SyntaxException {
		SequenceNode<BlockItemNode> root = unit.getRootNode();
		AST newAst;
		List<VariableDeclarationNode> inputVars = new ArrayList<>();
		List<BlockItemNode> newExternalList = new ArrayList<>();
		Map<String, VariableDeclarationNode> macroVars = new HashMap<>();

		OrdinaryEntity mainEntity = unit.getInternalOrExternalEntity("main");

		if (mainEntity == null) {
			throw new SyntaxException("missing main function", unit
					.getRootNode().getSource());
		}
		if (!(mainEntity instanceof Function)) {
			throw new SyntaxException("non-function entity with name \"main\"",
					mainEntity.getFirstDeclaration().getSource());
		}

		Function mainFunction = (Function) mainEntity;
		FunctionDefinitionNode mainDef = mainFunction.getDefinition();

		unit = renameStaticVariables(unit);
		unit.release();
		root = moveStaticVariables(root);
		processMalloc(root);
		// remove main prototypes...
		for (DeclarationNode decl : mainFunction.getDeclarations()) {
			if (!decl.isDefinition()) {
				decl.parent().removeChild(decl.childIndex());
			}
		}
		this.mainSource = mainDef.getSource();
		inputVars = processMainFunction(mainDef);
		processArgvRefs(mainDef.getBody());
		for (BlockItemNode inputVar : macroVars.values())
			newExternalList.add(inputVar);
		for (BlockItemNode inputVar : inputVars)
			newExternalList.add(inputVar);
		if (this.argcAssumption != null) {
			newExternalList.add(this.assumeFunctionDeclaration(argcAssumption
					.getSource()));
			newExternalList.add(argcAssumption);
		}
		// add my root
		newExternalList.add(this.myRootNode());
		for (BlockItemNode child : root) {
			if (child != null) {
				newExternalList.add(child);
				child.parent().removeChild(child.childIndex());
			}
		}
		root = nodeFactory.newSequenceNode(root.getSource(), "TranslationUnit",
				newExternalList);
		this.completeSources(root);
		newAst = astFactory.newAST(root, unit.getSourceFiles());
		return newAst;
	}

	private VariableDeclarationNode myRootNode() {
		return nodeFactory.newVariableDeclarationNode(mainSource,
				nodeFactory.newIdentifierNode(mainSource, GENERAL_ROOT),
				nodeFactory.newScopeTypeNode(mainSource),
				nodeFactory.newHereNode(mainSource));
	}

	private void processArgvRefs(ASTNode node) throws SyntaxException {
		if (node instanceof OperatorNode
				&& ((OperatorNode) node).getOperator() == Operator.SUBSCRIPT) {
			ASTNode parent = node.parent();

			if (!(parent instanceof OperatorNode && (((OperatorNode) parent)
					.getOperator() == Operator.ADDRESSOF))
					&& !(parent instanceof CastNode)) {
				OperatorNode subscript = (OperatorNode) node;
				ExpressionNode arg0 = subscript.getArgument(0);

				if (arg0.expressionKind() == ExpressionKind.IDENTIFIER_EXPRESSION) {
					IdentifierExpressionNode idExpr = (IdentifierExpressionNode) arg0;

					if (idExpr.getIdentifier().name().equals(this.argvName)) {
						OperatorNode newSubscript = subscript.copy();
						IdentifierExpressionNode newIdExpr = idExpr.copy();
						Source source = subscript.getSource();
						ExpressionNode addreessOf;

						newIdExpr.getIdentifier().setName(newArgvName);
						newSubscript.setChild(0, newIdExpr);
						newSubscript = nodeFactory.newOperatorNode(source,
								Operator.SUBSCRIPT, Arrays.asList(newSubscript,
										nodeFactory.newIntegerConstantNode(
												source, "0")));
						addreessOf = nodeFactory.newOperatorNode(source,
								Operator.ADDRESSOF,
								Arrays.asList((ExpressionNode) newSubscript));
						node.parent().setChild(node.childIndex(), addreessOf);
					}
				}
			}
		} else {
			for (ASTNode child : node.children())
				if (child != null)
					processArgvRefs(child);
		}
	}

	private void processMalloc(ASTNode node) {
		if (node instanceof FunctionCallNode) {
			FunctionCallNode funcCall = (FunctionCallNode) node;

			if (funcCall.getFunction().expressionKind() == ExpressionKind.IDENTIFIER_EXPRESSION) {
				IdentifierExpressionNode functionExpression = (IdentifierExpressionNode) funcCall
						.getFunction();
				String functionName = functionExpression.getIdentifier().name();

				if (functionName.equals(MALLOC)) {
					ASTNode parent = funcCall.parent();
					int callIndex = funcCall.childIndex();
					ExpressionNode myRootScope = this.identifierExpression(
							funcCall.getSource(), GENERAL_ROOT);
					ExpressionNode argument = funcCall.getArgument(0);

					functionExpression.getIdentifier().setName("$" + MALLOC);
					argument.parent().removeChild(argument.childIndex());
					funcCall.setArguments(nodeFactory.newSequenceNode(
							argument.getSource(), "Actual Parameters",
							Arrays.asList(myRootScope, argument)));
					if (!(parent instanceof CastNode)) {
						if (parent instanceof OperatorNode) {
							ExpressionNode lhs = ((OperatorNode) parent)
									.getArgument(0);
							Type type = lhs.getInitialType();
							TypeNode typeNode;
							CastNode castNode;

							if (type.kind() != TypeKind.POINTER)
								throw new CIVLSyntaxException(
										"The left hand side of a malloc call must be of pointer"
												+ " type.", lhs.getSource());
							typeNode = this.typeNode(lhs.getSource(), type);
							parent.removeChild(callIndex);
							castNode = nodeFactory.newCastNode(
									funcCall.getSource(), typeNode, funcCall);
							parent.setChild(callIndex, castNode);
						}
					}
				}
			}
		} else {
			for (ASTNode child : node.children()) {
				if (child != null)
					processMalloc(child);
			}
		}

	}

	/**
	 * Processes the original main function, including:
	 * <ul>
	 * <li>Removes all arguments of the function;</li>
	 * </ul>
	 * 
	 * @param mainFunction
	 *            The function definition node representing the original main
	 *            function.
	 * @param vars
	 *            The list of variable declaration nodes that are the arguments
	 *            of the original main function. These variables will be moved
	 *            up to the higher scope (i.e., the file scope of the final AST)
	 *            and become $input variables of the final AST. When invoking
	 *            this function, this parameter should be an empty list and this
	 *            function will update this list.
	 * @throws SyntaxException
	 */
	private List<VariableDeclarationNode> processMainFunction(
			FunctionDefinitionNode mainFunction) throws SyntaxException {
		List<VariableDeclarationNode> inputVars = new ArrayList<>();
		FunctionTypeNode functionType = mainFunction.getTypeNode();

		SequenceNode<VariableDeclarationNode> parameters = functionType
				.getParameters();
		int count = parameters.numChildren();

		if (count != 0 && count != 2) {
			if (count == 1) {
				if (parameters.getSequenceChild(0).getTypeNode().typeNodeKind() != TypeNodeKind.VOID)
					throw new SyntaxException(
							"The main function should have 0 or 2 parameters instead of "
									+ count, mainFunction.getSource());
			} else
				throw new SyntaxException(
						"The main function should have 0 or 2 parameters instead of "
								+ count, mainFunction.getSource());
		}
		if (count == 2) {
			VariableDeclarationNode argc = parameters.getSequenceChild(0);
			VariableDeclarationNode argv = parameters.getSequenceChild(1);
			VariableDeclarationNode CIVL_argc = argc.copy();
			VariableDeclarationNode CIVL_argv, _argv;
			String argcName = argc.getIdentifier().name();
			String argvName = argv.getIdentifier().name();
			String _argvName = "_" + argvName;
			String newArgcName = INPUT_PREFIX + argcName;
			Source source = argv.getTypeNode().getSource();
			TypeNode pointerOfPointerOfChar = nodeFactory.newPointerTypeNode(
					source,
					nodeFactory.newPointerTypeNode(source,
							this.basicType(BasicTypeKind.CHAR)));
			CompoundStatementNode functionBody = mainFunction.getBody();
			TypeNode arrayOfCharPointer = nodeFactory.newArrayTypeNode(
					source,
					nodeFactory.newPointerTypeNode(source,
							this.basicType(BasicTypeKind.CHAR)),
					nodeFactory.newIntegerConstantNode(source, MAX_ARGC));
			LoopNode forLoop;
			ForLoopInitializerNode loopInit;
			ExpressionNode condition, increment;
			StatementNode loopBody;
			ExpressionNode lhs, rhs;
			ExpressionNode assignArgv;

			this.argvName = argvName;
			this.newArgvName = INPUT_PREFIX + argvName;
			parameters.removeChild(0);
			parameters.removeChild(1);
			CIVL_argc.getTypeNode().setInputQualified(true);
			CIVL_argc.getIdentifier().setName(newArgcName);
			inputVars.add(CIVL_argc);
			CIVL_argv = inputArgvDeclaration(argv, newArgvName);
			inputVars.add(CIVL_argv);
			functionType.setParameters(nodeFactory.newSequenceNode(
					parameters.getSource(), "FormalParameterDeclarations",
					new ArrayList<VariableDeclarationNode>(0)));
			this.argcAssumption = this.argcAssumption(argc.getSource(),
					newArgcName);
			argc.setInitializer(this.identifierExpression(argc.getSource(),
					newArgcName));
			argv.setTypeNode(pointerOfPointerOfChar);
			_argv = argv.copy();
			_argv.getIdentifier().setName(_argvName);
			_argv.setTypeNode(arrayOfCharPointer);
			source = argv.getSource();
			loopInit = nodeFactory.newForLoopInitializerNode(
					source,
					Arrays.asList(this.variableDeclaration("i",
							this.basicType(BasicTypeKind.INT),
							nodeFactory.newIntegerConstantNode(source, "0"))));
			condition = nodeFactory.newOperatorNode(source, Operator.LT, Arrays
					.asList(this.identifierExpression(source, "i"), nodeFactory
							.newIntegerConstantNode(source, MAX_ARGC)));
			increment = nodeFactory.newOperatorNode(source,
					Operator.POSTINCREMENT,
					Arrays.asList(this.identifierExpression(source, "i")));
			// _argv[i]
			lhs = nodeFactory.newOperatorNode(source, Operator.SUBSCRIPT,
					Arrays.asList(this.identifierExpression(source, _argvName),
							this.identifierExpression(source, "i")));
			// CIVL_argv[i]
			rhs = nodeFactory.newOperatorNode(source, Operator.SUBSCRIPT,
					Arrays.asList(
							this.identifierExpression(source, newArgvName),
							this.identifierExpression(source, "i")));
			// CIVL_argv[i][0]
			rhs = nodeFactory.newOperatorNode(
					source,
					Operator.SUBSCRIPT,
					Arrays.asList(rhs,
							nodeFactory.newIntegerConstantNode(source, "0")));
			// &CIVL_argv[i][0]
			rhs = nodeFactory.newOperatorNode(source, Operator.ADDRESSOF,
					Arrays.asList(rhs));
			loopBody = nodeFactory.newExpressionStatementNode(nodeFactory
					.newOperatorNode(source, Operator.ASSIGN,
							Arrays.asList(lhs, rhs)));
			forLoop = nodeFactory.newForLoopNode(source, loopInit, condition,
					increment, loopBody, null);
			// _argv[0];
			assignArgv = nodeFactory.newOperatorNode(source,
					Operator.SUBSCRIPT, Arrays.asList(
							this.identifierExpression(source, _argvName),
							nodeFactory.newIntegerConstantNode(source, "0")));
			// &_argv[0];
			assignArgv = nodeFactory.newOperatorNode(source,
					Operator.ADDRESSOF, Arrays.asList(assignArgv));
			// argv = &_argv[0];
			assignArgv = nodeFactory.newOperatorNode(source, Operator.ASSIGN,
					Arrays.asList(this.identifierExpression(source, argvName),
							assignArgv));
			functionBody = this.addNodeToBeginning(functionBody,
					nodeFactory.newExpressionStatementNode(assignArgv));
			functionBody = this.addNodeToBeginning(functionBody, forLoop);
			functionBody = this.addNodeToBeginning(functionBody, argv);
			functionBody = this.addNodeToBeginning(functionBody, _argv);
			functionBody = this.addNodeToBeginning(functionBody, argc);
			mainFunction.setBody(functionBody);
		}
		return inputVars;
	}

	/**
	 * $assume 0 < argc && argc < MAX_ARGC;
	 * 
	 * @param source
	 * @param argcName
	 * @return
	 * @throws SyntaxException
	 */
	private ExpressionStatementNode argcAssumption(Source source,
			String argcName) throws SyntaxException {
		ExpressionNode lowerBound = nodeFactory.newOperatorNode(source,
				Operator.LT, Arrays.asList(
						nodeFactory.newIntegerConstantNode(source, "0"),
						this.identifierExpression(source, argcName)));
		ExpressionNode upperBound = nodeFactory.newOperatorNode(source,
				Operator.LT, Arrays.asList(
						this.identifierExpression(source, argcName),
						nodeFactory.newIntegerConstantNode(source, MAX_ARGC)));

		return nodeFactory.newExpressionStatementNode(this.functionCall(source,
				ASSUME, Arrays.asList((ExpressionNode) nodeFactory
						.newOperatorNode(source, Operator.LAND,
								Arrays.asList(lowerBound, upperBound)))));
	}

	/**
	 * Declares <code>$input char CIVL_argv[MAX_ARGC][];</code>.
	 * 
	 * @param oldArgv
	 * @return
	 * @throws SyntaxException
	 */
	private VariableDeclarationNode inputArgvDeclaration(
			VariableDeclarationNode oldArgv, String argvNewName)
			throws SyntaxException {
		VariableDeclarationNode __argv = oldArgv.copy();
		Source source = oldArgv.getSource();
		TypeNode arrayOfString = nodeFactory.newArrayTypeNode(
				source,
				nodeFactory.newArrayTypeNode(oldArgv.getSource(),
						this.basicType(BasicTypeKind.CHAR), null),
				nodeFactory.newIntegerConstantNode(source, MAX_ARGC));

		__argv.getIdentifier().setName(argvNewName);
		arrayOfString.setInputQualified(true);
		__argv.setTypeNode(arrayOfString);
		return __argv;
	}

	public enum ArgvTypeKind {
		POINTER_POINTER_CHAR, ARRAY_POINTER_CHAR, ARRAY_ARRAY_CAHR
	};

	// private ArgvTypeKind analyzeArgvType(TypeNode type) throws
	// SyntaxException {
	// TypeNodeKind typeKind = type.typeNodeKind();
	//
	// switch (typeKind) {
	// case POINTER:
	// return ArgvTypeKind.POINTER_POINTER_CHAR;
	// case ARRAY:
	// ArrayTypeNode arrayType = (ArrayTypeNode) type;
	// TypeNode arrayEleType = arrayType.getElementType();
	// TypeKind arrayEleTypeKind = arrayEleType.getType().kind();
	//
	// if (arrayEleTypeKind == TypeKind.POINTER)
	// return ArgvTypeKind.ARRAY_POINTER_CHAR;
	// else if (arrayEleTypeKind == TypeKind.ARRAY)
	// return ArgvTypeKind.ARRAY_ARRAY_CAHR;
	// default:
	// throw new SyntaxException("Invalid type " + type.getType()
	// + " for argv of main.", null);
	// }
	// }

	private CompoundStatementNode addNodeToBeginning(
			CompoundStatementNode compoundNode, BlockItemNode node) {
		int numChildren = compoundNode.numChildren();
		List<BlockItemNode> nodeList = new ArrayList<>(numChildren + 1);

		nodeList.add(node);
		for (int i = 0; i < numChildren; i++) {
			BlockItemNode child = compoundNode.getSequenceChild(i);

			nodeList.add(child);
			compoundNode.removeChild(i);
		}
		return nodeFactory.newCompoundStatementNode(compoundNode.getSource(),
				nodeList);
	}

	private AST renameStaticVariables(AST ast) throws SyntaxException {
		Map<Entity, String> newNameMap = new HashMap<>();
		NameTransformer staticVariableNameTransformer;

		newNameMap = newNameMapOfStaticVariables(ast.getRootNode(), newNameMap);
		staticVariableNameTransformer = Transform.nameTransformer(newNameMap,
				astFactory);
		return staticVariableNameTransformer.transform(ast);
	}

	// TODO can you have static for function parameters?
	// TODO what if the initializer of the variable node access some variables
	// not declared in the root scope?
	/**
	 * Computes the new name map of static variables. A static variable "var" is
	 * renamed to "var$n", where n is the current static variable ID.
	 * 
	 * @param node
	 * @param newNames
	 * @return
	 */
	private Map<Entity, String> newNameMapOfStaticVariables(ASTNode node,
			Map<Entity, String> newNames) {
		if (node instanceof VariableDeclarationNode) {
			VariableDeclarationNode variable = (VariableDeclarationNode) node;

			if (variable.hasStaticStorage()) {
				String oldName = variable.getName();
				String newName = oldName + separator + this.static_var_count++;

				newNames.put(variable.getEntity(), newName);
				this.static_variables.add(variable);
			}
		} else {
			for (ASTNode child : node.children()) {
				if (child == null)
					continue;
				newNames = newNameMapOfStaticVariables(child, newNames);
			}
		}
		return newNames;
	}

	private SequenceNode<BlockItemNode> moveStaticVariables(
			SequenceNode<BlockItemNode> root) {
		if (this.static_variables.size() < 1)
			return root;

		List<BlockItemNode> newChildren = new LinkedList<>();
		int count = root.numChildren();

		for (VariableDeclarationNode var : this.static_variables) {
			var.remove();
			newChildren.add(var);
		}
		for (int i = 0; i < count; i++) {
			BlockItemNode child = root.getSequenceChild(i);

			if (child == null)
				continue;
			child.remove();
			newChildren.add(child);
		}
		return nodeFactory
				.newTranslationUnitNode(root.getSource(), newChildren);
	}

}