package cd.transform.analysis;

import cd.ir.BasicBlock;
import cd.ir.ControlFlowGraph;

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

import static java.util.Collections.unmodifiableMap;

/**
 * The abstract superclass of all data-flow analyses. This class provides a framework to
 * implement concrete analyses by providing {@link #initialState()},
 * {@link #startState()}, {@link #transferFunction(BasicBlock, Object)}, and
 * {@link #join(Set)} methods.
 *
 * @param <State> The type of states the analysis computes, specified by a concrete subclass.
 *                Typically, this is a set or map type.
 */
public abstract class DataFlowAnalysis<State> {

	protected final ControlFlowGraph cfg;
	Map<BasicBlock, State> inStates;
	Map<BasicBlock, State> outStates;

	public DataFlowAnalysis(ControlFlowGraph cfg) {
		this.cfg = cfg;
	}

	/**
	 * Returns the in-state of basic block <code>block</code>.
	 */
	public State inStateOf(BasicBlock block) {
		return inStates.getOrDefault(block, startState());
	}

	/**
	 * Returns the out-state of basic block <code>block</code>.
	 */
	public State outStateOf(BasicBlock block) {
		return outStates.get(block);
	}

	protected void emptyStates() {
		inStates = new HashMap<>();
		outStates = new HashMap<>();
	}

	/**
	 * Do forward flow fixed-point iteration until out-states do not change anymore.
	 * Subclasses should call this method in their constructor after the required
	 * initialization.
	 */
	protected void iterate() {
		inStates = new HashMap<>();
		outStates = new HashMap<>();
		for (BasicBlock block : cfg.allBlocks)
			outStates.put(block, initialState());

		Set<BasicBlock> todo = new HashSet<>();
		todo.addAll(cfg.allBlocks);
		while (!todo.isEmpty()) {
			BasicBlock block = todo.iterator().next();
			todo.remove(block);

			/* calculate in-state */
			State inState;
			if (block == cfg.start)
				inState = startState();
			else {
				Set<State> predOutStates = new HashSet<>();
				for (BasicBlock pred : block.predecessors)
					predOutStates.add(outStates.get(pred));
				inState = join(predOutStates);
			}
			inStates.put(block, inState);

			State newOutState = transferFunction(block, inState);

			/* if out-state changed, recalculate successors */
			if (!newOutState.equals(outStates.get(block))) {
				outStates.put(block, newOutState);
				todo.addAll(block.successors);
			}
		}
		outStates = unmodifiableMap(outStates);
	}

	/**
	 * Returns the initial state for all blocks except the {@link ControlFlowGraph#start start}
	 * block.
	 */
	protected abstract State initialState();

	/**
	 * Returns the initial state for the {@link ControlFlowGraph#start start} block.
	 */
	protected abstract State startState();

	/**
	 * Calculates the out-state for a basic block <code>block</code> and an in-state
	 * <code>inState</code>
	 */
	protected abstract State transferFunction(BasicBlock block, State inState);

	/**
	 * Merges together several out-states and returns the in-state for the transfer function.
	 */
	protected abstract State join(Set<State> states);

	/**
	 * DataFlowAnalysis implementation with reversed direction. The methods are overridden in
	 * order to reflect the change in the parameters' names and in the javadoc.
	 */
	public static abstract class Reverse<A> extends DataFlowAnalysis<A> {
		public Reverse(ControlFlowGraph cfg) {
			super(cfg);
		}

		/**
		 * Do backward flow fixed-point iteration until in-states do not change anymore.
		 * Subclasses should call this method in their constructor after the required
		 * initialization.
		 */
		@Override
		protected void iterate() {
			inStates = new HashMap<>();
			outStates = new HashMap<>();
			for (BasicBlock block : cfg.allBlocks)
				inStates.put(block, initialState());

			Set<BasicBlock> todo = new HashSet<>();
			todo.addAll(cfg.allBlocks);
			while (!todo.isEmpty()) {
				BasicBlock block = todo.iterator().next();
				todo.remove(block);

				/* calculate out-state */
				A outState;
				if (block == cfg.end)
					outState = startState();
				else {
					Set<A> predInStates = new HashSet<>();
					for (BasicBlock pred : block.successors)
						predInStates.add(inStates.get(pred));
					outState = join(predInStates);
				}
				outStates.put(block, outState);

				A newInState = transferFunction(block, outState);

				/* if in-state changed, recalculate successors */
				if (!newInState.equals(inStates.get(block))) {
					inStates.put(block, newInState);
					todo.addAll(block.predecessors);
				}
			}
			inStates = unmodifiableMap(inStates);
		}

		/**
		 * Returns the initial state for all blocks except the {@link ControlFlowGraph#end end}
		 * block.
		 */
		@Override
		protected abstract A initialState();

		/**
		 * Returns the initial state for the {@link ControlFlowGraph#end end} block.
		 */
		@Override
		protected abstract A startState();

		/**
		 * Calculates the in-state for a basic block <code>block</code> and an out-state
		 * <code>outState</code>
		 */
		@Override
		protected abstract A transferFunction(BasicBlock block, A outState);

		/**
		 * Merges together several in-states and returns the out-state for the transfer function.
		 */
		@Override
		protected abstract A join(Set<A> states);
	}
}
