package cd.transform.analysis;

import cd.ir.Ast;
import cd.ir.Ast.Assign;
import cd.ir.AstVisitor;
import cd.ir.BasicBlock;
import cd.ir.ControlFlowGraph;
import cd.ir.Symbol.VariableSymbol;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

public class LiveVarAnalysis extends DataFlowAnalysis.Reverse<Set<VariableSymbol>> {
	private Map<BasicBlock, Set<VariableSymbol>> killSets = new HashMap<>();
	private Map<BasicBlock, Set<VariableSymbol>> genSets = new HashMap<>();

	public LiveVarAnalysis(ControlFlowGraph cfg) {
		super(cfg);

		// For faster iteration, the transfer function will just use kill and gen sets
		for (BasicBlock block : cfg.allBlocks)
			killGenSetOf(block);

		iterate();
	}

	private void killGenSetOf(BasicBlock block) {
		Set<VariableSymbol> kill = new HashSet<>();
		Set<VariableSymbol> gen = new HashSet<>();

		if (block.condition != null)
			new Visitor().visit(block.condition, gen);
		for (int i = block.stmts.size() - 1; i >= 0; i--) {
			VariableSymbol var = new Visitor().visit(block.stmts.get(i), gen);
			if (var != null) {
				// TODO: review this code
				kill.add(var);
				gen.remove(var);
			}
		}

		killSets.put(block, kill);
		genSets.put(block, gen);
	}

	@Override
	protected Set<VariableSymbol> initialState() {
		return new HashSet<>();
	}

	@Override
	protected Set<VariableSymbol> startState() {
		return new HashSet<>();
	}

	@Override
	protected Set<VariableSymbol> transferFunction(BasicBlock block, Set<VariableSymbol> inState) {
		Set<VariableSymbol> result = new HashSet<>(inState);
		result.removeAll(killSets.get(block));
		result.addAll(genSets.get(block));
		return result;
	}

	@Override
	protected Set<VariableSymbol> join(Set<Set<VariableSymbol>> states) {
		Set<VariableSymbol> result = new HashSet<>();
		states.forEach(result::addAll);
		return result;
	}

	public static class Visitor extends AstVisitor<VariableSymbol, Set<VariableSymbol>> {
		@Override
		public VariableSymbol assign(Assign ast, Set<VariableSymbol> arg) {
			visit(ast.right(), arg);
			if (ast.left() instanceof Ast.Var)
				return visit(ast.left(), arg);
			else
				return null;
		}

		@Override
		public VariableSymbol var(Ast.Var ast, Set<VariableSymbol> arg) {
			if (!ast.sym.type.isReferenceType()) {
				arg.add(ast.sym);
				return ast.sym;
			}
			return null;
		}
	}
}
