package cd.transform.optimizers;

import cd.ir.Ast;
import cd.ir.Ast.BooleanConst;
import cd.ir.Ast.Expr;
import cd.ir.Ast.IntConst;
import cd.ir.Ast.Var;
import cd.ir.BasicBlock;
import cd.ir.CompileTimeEvaluator;
import cd.ir.ControlFlowGraph;
import cd.ir.Symbol.PrimitiveTypeSymbol;
import cd.ir.Symbol.VariableSymbol;
import cd.transform.analysis.ConstantAnalysis;

import java.util.Map;
import java.util.Optional;

import static cd.Config.TRUE;

/**
 * Class that propagates constant values to avoid accessing variables at runtime,
 * and also includes simplifications to arithmetic operations.
 */
public class ConstantPropagator extends OptimizerVisitor<Map<VariableSymbol, Integer>> {

	private CompileTimeEvaluator cte = new CompileTimeEvaluator();

	@Override
	public boolean go(ControlFlowGraph cfg) {
		setChanged(false);

		ConstantAnalysis constantAnalysis = new ConstantAnalysis(cfg);
		for (BasicBlock block : cfg.allBlocks) {
			block.accept(this, constantAnalysis.inStateConstantsOf(block));
		}

		return hasChanged();
	}

	@Override
	public Ast assign(Ast.Assign ast, Map<VariableSymbol, Integer> arg) {
		// Visit and replace only the RHS of the statement
		Expr lhs = ast.left();
		ast.setLeft(null);
		visitChildren(ast, arg);
		ast.setLeft(lhs);

		// Update the map (only with constant values)
		if (ast.left() instanceof Var && !ast.left().type.isReferenceType()) {
			Optional<Integer> rightValue = cte.calc(ast.right(), arg);
			if (rightValue.isPresent())
				arg.put(((Var) ast.left()).sym, rightValue.get());
			else
				arg.remove(((Var) ast.left()).sym);
		}
		return ast;
	}

	@Override
	public Ast var(Var ast, Map<VariableSymbol, Integer> arg) {
		Integer value = arg.get(ast.sym);
		if (value == null)
			return ast;
		if (ast.type == PrimitiveTypeSymbol.intType)
			return new IntConst(value);

		assert ast.type == PrimitiveTypeSymbol.booleanType;
		return new BooleanConst(value == TRUE);
	}
}
