package cd.transform.analysis;

import cd.ir.*;
import cd.ir.Ast.Var;
import cd.ir.Symbol.VariableSymbol;

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

public abstract class MaybeFlowAnalysis<E> extends DataFlowAnalysis<Map<VariableSymbol, MaybeC<E>>> {
	protected final CompileTimeEvaluator cte = new CompileTimeEvaluator();
	private final Set<VariableSymbol> allVars = new HashSet<>();

	public MaybeFlowAnalysis(ControlFlowGraph cfg, boolean primitiveVars, boolean referenceVars) {
		super(cfg);

		// Obtain all variables
		for (BasicBlock block : cfg.allBlocks) {
			VarVisitor visitor = new VarVisitor(primitiveVars, referenceVars);
			block.accept(visitor, allVars);
		}

		if (!allVars.isEmpty())
			iterate();
		else
			emptyStates();
	}

	public Map<VariableSymbol, E> inStateConstantsOf(BasicBlock block) {
		return getConstants(inStateOf(block));
	}

	public Map<VariableSymbol, E> outStateConstantsOf(BasicBlock block) {
		return getConstants(outStateOf(block));
	}

	private Map<VariableSymbol, E> getConstants(Map<VariableSymbol, MaybeC<E>> map) {
		Map<VariableSymbol, E> result = new HashMap<>();
		if (map == null) return result;
		map.forEach((key, value) -> {
			if (value.isConstant())
				result.put(key, value.getValue());
		});
		return result;
	}

	/**
	 * Reads a statement and updates the current state accordingly,
	 */
	public abstract void transferStmt(Ast.Stmt stmt, Map<VariableSymbol, MaybeC<E>> state);

	/**
	 * Default value for each constant at the beginning of the start block
	 */
	protected abstract MaybeC<E> defaultInit(VariableSymbol varSym);

	@Override
	protected Map<VariableSymbol, MaybeC<E>> initialState() {
		Map<VariableSymbol, MaybeC<E>> map = new HashMap<>();
		allVars.forEach(v -> map.put(v, MaybeC.anyValue()));
		return map;
	}

	@Override
	protected Map<VariableSymbol, MaybeC<E>> startState() {
		Map<VariableSymbol, MaybeC<E>> map = new HashMap<>();
		allVars.forEach(v -> map.put(v, defaultInit(v)));
		return map;
	}

	@Override
	protected Map<VariableSymbol, MaybeC<E>> transferFunction(BasicBlock block, Map<VariableSymbol, MaybeC<E>> inState) {
		Map<VariableSymbol, MaybeC<E>> outState = new HashMap<>();
		inState.forEach((key, value) -> outState.put(key, new MaybeC<>(value)));
		for (Ast.Stmt stmt : block.stmts)
			transferStmt(stmt, outState);
		return outState;
	}

	@Override
	protected Map<VariableSymbol, MaybeC<E>> join(Set<Map<VariableSymbol, MaybeC<E>>> states) {
		// The basic condition is that the value must be the same in all maps for the constant to be included
		Map<VariableSymbol, MaybeC<E>> result = new HashMap<>();
		for (VariableSymbol sym : allVars) {
			MaybeC<E> value = MaybeC.anyValue();
			for (Map<VariableSymbol, MaybeC<E>> map : states) {
				value.joinWith(map.get(sym));
				if (value.isNotConstant())
					break; // It cannot change once it is notConstant
			}
			result.put(sym, value);
		}
		return result;
	}

	class VarVisitor extends AstVisitor<Void, Set<VariableSymbol>> {
		private final boolean primitive, reference;

		VarVisitor(boolean primitive, boolean reference) {
			this.primitive = primitive;
			this.reference = reference;
		}

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