package cd.transform.optimizers;

import cd.ir.Ast;
import cd.ir.Ast.*;
import cd.ir.BasicBlock;
import cd.ir.CompileTimeEvaluator;
import cd.ir.ControlFlowGraph;

import java.util.Optional;

import static cd.Config.TRUE;
import static cd.ir.Ast.UnaryOp.UOp.U_MINUS;

public class TrivialOpRemoval extends OptimizerVisitor<Void> {
	private CompileTimeEvaluator cte = new CompileTimeEvaluator();

	@Override
	public boolean go(ControlFlowGraph cfg) {
		setChanged(false);
		for (BasicBlock block : cfg.allBlocks)
			block.accept(this, null);
		return hasChanged();
	}

	@Override
	public Ast binaryOp(BinaryOp ast, Void arg) {
		visitChildren(ast, arg);
		Optional<Integer> rightRes = cte.calc(ast.right());
		Optional<Integer> leftRes = cte.calc(ast.left());

		if (leftRes.isPresent() && rightRes.isPresent()) {
			// Skip the operation if both arguments' values are known
			Optional<Integer> resOpt = cte.calc(ast);
			if (!resOpt.isPresent())
				return ast;

			switch (ast.operator) {
				case B_TIMES:
				case B_DIV:
				case B_MOD:
				case B_PLUS:
				case B_MINUS:
					return new IntConst(resOpt.get());
				default:
					return new BooleanConst(resOpt.get() == TRUE);
			}

		} else if (!leftRes.isPresent() && !rightRes.isPresent()) {
			// Nothing can be done if no argument is known
			return ast;
		}

		switch (ast.operator) {
			case B_PLUS: // a + 0 = a (commutative)
				if (leftRes.isPresent() && leftRes.get() == 0)
					return ast.right();
				else if (rightRes.isPresent() && rightRes.get() == 0)
					return ast.left();
				break;
			case B_MINUS: // a - 0 = a | 0 - a = -a
				if (leftRes.isPresent() && leftRes.get() == 0)
					return new UnaryOp(U_MINUS, ast.right());
				else if (rightRes.isPresent() && rightRes.get() == 0)
					return ast.left();
				break;
			case B_TIMES: // a * 0 = 0 | a * 1 = a | a * -1 = -a (commutative)
				if (rightRes.isPresent()) {
					if (rightRes.get() == 1)
						return ast.left();
					else if (rightRes.get() == -1)
						return new UnaryOp(U_MINUS, ast.left());
//				else if (rightRes.get() == 0)
//					return new IntConst(0);
				} else {
					if (leftRes.get() == 1)
						return ast.right();
					else if (leftRes.get() == -1)
						return new UnaryOp(U_MINUS, ast.right());
//				else if (leftRes.get() == 0)
//					return new IntConst(0);
				}
				break;
			case B_DIV: // 0 / a = 0 | a / 1 = a | a / -1 = -a
				if (rightRes.isPresent()) {
					if (rightRes.get() == 1)
						return ast.left();
					else if (rightRes.get() == -1)
						return new UnaryOp(U_MINUS, ast.left());
//				} else if (leftRes.get() == 0)
//					return new IntConst(0);
				}
				break;
			case B_MOD: // 0 % a = 0 | a % 1 = 0
//			if (rightValid && rightValue == 1)
//				return new IntConst(0);
//			else if (leftValid && leftValue == 0)
//				return new IntConst(0);
				break;
		}
		return ast;
	}

	@Override
	public Ast unaryOp(UnaryOp ast, Void arg) {
		visitChildren(ast, arg);
		Optional<Integer> valueOpt = cte.calc(ast);

		if (valueOpt.isPresent()) {
			// Skip the operation if the argument's value is known
			switch (ast.operator) {
				case U_BOOL_NOT:
					return new BooleanConst(valueOpt.get() == TRUE);
				default:
					return new IntConst(valueOpt.get());
			}
		} else {
			switch (ast.operator) {
				case U_PLUS:
					// The plus operator is meaningless
					return ast.arg();
				default:
					// Nested unary operators with the same operator are meaningless
					if (ast.arg() instanceof UnaryOp) {
						UnaryOp inner = (UnaryOp) ast.arg();
						if (inner.operator == ast.operator)
							return inner.arg();
						else return ast;
					} else {
						return ast;
					}
			}
		}
	}


	@Override
	public Ast cast(Cast ast, Void arg) {
		visitChildren(ast, arg);
		if (ast.type == ast.arg().type)
			return ast.arg();
		return ast;
	}
}
