/**
 * Copyright (c) 2005 by Kevin Bierhoff
 */
package edu.cmu.cs.l3.trace;

import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedList;
import java.util.Vector;

import edu.cmu.cs.l3.canon.Block;
import edu.cmu.cs.l3.frame.Frame;
import edu.cmu.cs.l3.temp.Label;
import edu.cmu.cs.l3.temp.Temp;
import edu.cmu.cs.l3.tree.IRBinop;
import edu.cmu.cs.l3.tree.IRCJump;
import edu.cmu.cs.l3.tree.IRCall;
import edu.cmu.cs.l3.tree.IRConst;
import edu.cmu.cs.l3.tree.IRExpressionList;
import edu.cmu.cs.l3.tree.IRJump;
import edu.cmu.cs.l3.tree.IRLabel;
import edu.cmu.cs.l3.tree.IRMove;
import edu.cmu.cs.l3.tree.IRName;
import edu.cmu.cs.l3.tree.IRStatementExpression;
import edu.cmu.cs.l3.tree.IRStatementList;
import edu.cmu.cs.l3.tree.IRTemp;

/**
 * @author Kevin
 *
 */
public class TraceInstrumentation {

	private Vector blocks;
	private PathGraph pg;
	private PathGraph completeGraph;
	private HashMap instrumentation = new HashMap();
	private Temp pathTemp;
	private Frame frame;
	private int counterIndex;
	private Block uniqueExit;

	/**
	 * Initializes trace instrumentation from a given set of blocks.
	 * @param blocks
	 * @param frame
	 * @param counterIndex
	 */
	public TraceInstrumentation(Vector blocks, Frame frame, int counterIndex) {
		super();
		this.blocks = blocks;
		this.pg = new PathGraph(blocks);
		this.completeGraph = new PathGraph(blocks);
		this.frame = frame;
		this.pathTemp = new Temp();
		this.counterIndex = counterIndex;
		
		computeInstrumentation();
	}
	
	/**
	 * Computes trace information for a given path graph.
	 * @param pg
	 * @param frame
	 * @param counterIndex
	 */
	public TraceInstrumentation(PathGraph pg, Frame frame, int counterIndex) {
		super();
		this.blocks = pg.getBlocks();
		this.pg = pg;
		this.completeGraph = new PathGraph(this.blocks);
		this.frame = frame;
		this.pathTemp = new Temp();
		this.counterIndex = counterIndex;
		
		computeInstrumentation();
	}
	
	public void placeInstrumentation() {
		for(Iterator it = pg.getEdges().iterator(); it.hasNext(); ) {
			PathGraphEdge e = (PathGraphEdge) it.next();
			if(isInstrumented(e) == false)
				continue;
				
			Instrumentation i = getInstrumentation(e);
			System.out.println("Instrument " + e + " with " + i);
			
			if(e.isProgramEdge()) {
				Block b = breakEdge(e.getOrigin(), e.getTarget());
				placeInstrumentation(b, i);
			}
			else if(e.isSkipEdge()) {
				Block b = e.getOrigin().getBlock();
				placeInstrumentation(b, i);
			}
			else if(e.isCycleEdge()) {
				Block b = pg.getExitVertex().getBlock();
				placeInstrumentation(b, i);
			}
			// handle loop edges separately
		}
		// handle loop edges
		for(Iterator it = pg.getBackEdges().iterator(); it.hasNext(); ) {
			BackEdge e = (BackEdge) it.next();
			PathGraphEdge dummyIn = e.getDummyIn();
			PathGraphEdge dummyOut = e.getDummyOut();
			if(isInstrumented(dummyIn) == false && isInstrumented(dummyOut) == false)
				continue;
			Block b = breakEdge(e.getOrigin(), e.getTarget());
			// do after-loop instrumentation first (increase counter)
			if(isInstrumented(dummyOut)) {
				// dummyOut != null if true
				placeInstrumentation(b, getInstrumentation(dummyOut));
			}
			// TODO else increment counter anyway ??
			// before-loop instrumentation (initialize path register)
			if(isInstrumented(dummyIn)) {
				// dummyIn != null always
				placeInstrumentation(b, getInstrumentation(dummyIn));
			}
			// TODO else set path register to 0 ??
		}
	}
	
	public Block uniqueEntry() {
		return pg.getEntryVertex().getBlock();
	}
	
	public Block uniqueExit() {
		if(uniqueExit != null) return uniqueExit;
		
		// TODO find a unique exit node without using path graph
		int realExitCount = 0;
		Block result = null;
		for(Iterator it = completeGraph.getExitVertex().getIncomingEdges().iterator(); it.hasNext(); ) {
			PathGraphEdge e = (PathGraphEdge) it.next();
			if(e.isProgramEdge()) {
				++realExitCount;
				result = e.getOrigin().getBlock();
			}
		}
		if(realExitCount == 1) {
			uniqueExit = result;
			return result;
		}
		
		Label uniqueExitLab = new Label();
		result = new Block(new IRLabel(uniqueExitLab), 
				new IRStatementList(new IRJump(new IRLabel(frame.getReturnLabel()))));
		blocks.add(result);
		for(Iterator it = completeGraph.getExitVertex().getIncomingEdges().iterator(); it.hasNext(); ) {
			PathGraphEdge e = (PathGraphEdge) it.next();
			if(e.isProgramEdge()) {
				Block b = e.getOrigin().getBlock();
				((IRJump) b.getJump()).setTarget(new IRLabel(uniqueExitLab));
			}
		}
		uniqueExit = result;
		return result;
	}
	
	private void placeInstrumentation(Block b, Instrumentation i) {
		// if trying to instrument exit block, instrument pre-exits instead
		if(b == pg.getExitVertex().getBlock()) {
			System.out.println("Instrumenting exit vertex. This normally shouldn't happen.");
			b = uniqueExit();
//			for(Iterator it = pg.getPreExitVertices().iterator(); it.hasNext(); ) {
//				placeInstrumentation(((PathGraphVertex) it.next()).getBlock(), i);
//			}
		}
		
		// see if path register gets initialized or incremented
		if(i.isInitPathRegister()) {
			IRMove initStmt = new IRMove(new IRTemp(pathTemp), new IRConst(i.getPathRegisterIncrement()));
			b.getBody().insertStatement(b.getBody().getSize() - 1, initStmt);
		}
		else if(i.isIncPathRegister() && !i.isInitPathRegister()) {
			IRBinop addExpr = new IRBinop(IRBinop.PLUS, new IRTemp(pathTemp), new IRConst(i.getPathRegisterIncrement()));
			IRMove addStmt = new IRMove(new IRTemp(pathTemp), addExpr);
			b.getBody().insertStatement(b.getBody().getSize() - 1, addStmt);
		}
		
		// increment counter if necessary
		if(i.isIncCounter()) {
			/*IRMove addrStmt = new IRMove(new IRTemp(addrComputationTemp),
					new IRBinop(IRBinop.PLUS, new IRTemp(countAddrTemp), 
							new IRBinop(IRBinop.MUL, new IRTemp(pathTemp), new IRConst(getCounterWidth()))));
			b.getBody().insertStatement(b.getBody().getSize() - 1, addrStmt);
			IRMove loadCounter = new IRMove(new IRTemp(counterIncTemp), new IRMem(new IRTemp(addrComputationTemp)));
			b.getBody().insertStatement(b.getBody().getSize() - 1, loadCounter);
			IRBinop incExpr = new IRBinop(IRBinop.PLUS, new IRTemp(counterIncTemp), new IRConst(1));
			IRMove incStmt = new IRMove(new IRMem(new IRTemp(addrComputationTemp)), incExpr);*/
			IRExpressionList incCallArgs = new IRExpressionList(new IRConst(counterIndex), 
					new IRExpressionList(new IRTemp(pathTemp)));
			IRCall incCall = new IRCall(new IRName("_inc_count"), incCallArgs);
			IRStatementExpression incStmt = new IRStatementExpression(incCall);
			b.getBody().insertStatement(b.getBody().getSize() - 1, incStmt);
		}
	}

	/**
	 * @return
	 */
	private int getCounterWidth() {
		return frame.getWordSize();
	}

	/**
	 * @param e
	 * @return
	 */
	private Block breakEdge(PathGraphVertex origin, PathGraphVertex target) {
		Block b = origin.getBlock();
		// origin has only one successor, so we can just use that block
		if(b.getJump() instanceof IRJump)
			return b;

		// label for new block
		Label lab = new Label();
		// statement list that contains only a jump to b
		IRStatementList justTheJump = new IRStatementList(new IRJump(new IRLabel(target.getBlock().getLabel().getLabel())));
		// the new block
		Block preHeader = new Block(new IRLabel(lab), justTheJump);
		// we append it to the end of the list of blocks
		blocks.add(preHeader);

		String targetName = target.getBlock().getLabel().toString();
		IRCJump cjump = (IRCJump) b.getJump();
		if(cjump.getFalseLabel().toString().equals(targetName)) {
			cjump.setFalseLabel(new IRLabel(lab));
		}
		if(cjump.getTrueLabel().toString().equals(targetName)) {
			cjump.setTrueLabel(new IRLabel(lab));
		}
		return preHeader;
	}

	/**
	 * Imlements Ball & Larus' algorithm for efficiently placing instrumentation
	 * based on chord increments computed in {@link PathGraph}.
	 */
	protected void computeInstrumentation() {
		LinkedList workingSet;
		
		// register initialization code
		workingSet = new LinkedList();
		workingSet.addLast(pg.getEntryVertex());
		while(workingSet.isEmpty() == false) {
			PathGraphVertex v = (PathGraphVertex) workingSet.removeFirst();
			for(Iterator it = v.getOutgoingEdges().iterator(); it.hasNext(); ) {
				PathGraphEdge e = (PathGraphEdge) it.next();
				if(e.isSpanning() == false) {
					instrumentInitRegister(e, e.getIncrement());
				}
				else if(e.getTarget().getIncomingEdges().size() == 1) {
					workingSet.addLast(e.getTarget());
				}
				else {
					instrumentInitRegister(e, 0);
				}
			}
		}
		
		// memory increment code
		workingSet = new LinkedList();
		workingSet.addLast(pg.getExitVertex());
		while(workingSet.isEmpty() == false) {
			PathGraphVertex w = (PathGraphVertex) workingSet.removeFirst();
			for(Iterator it = w.getIncomingEdges().iterator(); it.hasNext(); ) {
				PathGraphEdge e = (PathGraphEdge) it.next();
				if(e.isSpanning() == false) {
					instrumentIncCounter(e, e.getIncrement());
				}
				else if(e.getOrigin().getOutgoingEdges().size() == 1) {
					workingSet.addLast(e.getOrigin());
				}
				else {
					instrumentIncCounter(e, 0);
				}
			}
		}
		
		// register increment code
		for(Iterator it = pg.getEdges().iterator(); it.hasNext(); ) {
			PathGraphEdge e = (PathGraphEdge) it.next();
			if(e.isSpanning() == false && isInstrumented(e) == false) {
				instrumentIncRegister(e, e.getIncrement());
			}
		}
	}

	/**
	 * @param e
	 * @param increment
	 */
	private void instrumentIncRegister(PathGraphEdge e, int increment) {
		Instrumentation i = instrument(e);
		i.setIncPathRegister(true);
		i.setPathRegisterIncrement(increment);
	}

	/**
	 * @param e
	 * @param increment
	 */
	private void instrumentIncCounter(PathGraphEdge e, int increment) {
		Instrumentation i = instrument(e);
		i.setIncCounter(true);
		if(increment != 0) {
			i.setIncPathRegister(true);
			i.setPathRegisterIncrement(increment);
		}
	}

	/**
	 * @param e
	 * @param increment
	 */
	private void instrumentInitRegister(PathGraphEdge e, int increment) {
		Instrumentation i = instrument(e);
		i.setInitPathRegister(true);
		i.setPathRegisterIncrement(increment);
	}

	/**
	 * @param e
	 * @return
	 */
	public Instrumentation getInstrumentation(PathGraphEdge e) {
		if(instrumentation.containsKey(e))
			return (Instrumentation) instrumentation.get(e);
		return null;
	}

	/**
	 * @param e
	 * @return
	 */
	protected boolean isInstrumented(PathGraphEdge edge) {
		return edge != null && instrumentation.containsKey(edge);
	}

	/**
	 * @param edge
	 * @param instr
	 */
	protected Instrumentation instrument(PathGraphEdge edge) {
		Instrumentation edgeInstr;
		if(instrumentation.containsKey(edge)) {
			edgeInstr = (Instrumentation) instrumentation.get(edge);
		}
		else {
			edgeInstr = new Instrumentation();
			instrumentation.put(edge, edgeInstr);
		}
		return edgeInstr;
	}

	/**
	 * @return
	 */
	public int getPathCount() {
		return pg.getPathCount();
	}

}
