// Mike Maxim
// IR canonicalizer (taken mostly from Appel)

package edu.cmu.cs.l3.canon;

import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.FileWriter;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.io.PrintWriter;
import java.util.HashMap;
import java.util.Vector;

import edu.cmu.cs.l3.frame.Frame;
import edu.cmu.cs.l3.temp.Temp;
import edu.cmu.cs.l3.trace.PathGraph;
import edu.cmu.cs.l3.trace.TraceAnalysis;
import edu.cmu.cs.l3.trace.TraceInstrumentation;
import edu.cmu.cs.l3.tree.IRCall;
import edu.cmu.cs.l3.tree.IRConst;
import edu.cmu.cs.l3.tree.IRESeq;
import edu.cmu.cs.l3.tree.IRExpression;
import edu.cmu.cs.l3.tree.IRExpressionList;
import edu.cmu.cs.l3.tree.IRMove;
import edu.cmu.cs.l3.tree.IRSeq;
import edu.cmu.cs.l3.tree.IRStatement;
import edu.cmu.cs.l3.tree.IRStatementExpression;
import edu.cmu.cs.l3.tree.IRStatementList;
import edu.cmu.cs.l3.tree.IRTemp;

class MoveCall extends IRStatement {
    IRTemp dst;
    IRCall src;
    MoveCall(IRTemp d, IRCall s) {dst=d; src=s;}
    public IRExpressionList getComponents() {return src.getComponents();}
    public IRStatement build(IRExpressionList kids) {
	return new IRMove(dst, src.build(kids));
    }
}

class ExpCall extends IRStatement {
    IRCall call;
    ExpCall(IRCall c) {call=c;}
    public IRExpressionList getComponents() {return call.getComponents();}
    public IRStatement build(IRExpressionList kids) {
	return new IRStatementExpression(call.build(kids));
    }
}

class StatementExpressionList {

    public StatementExpressionList(IRStatement s, IRExpressionList e) {
	m_stm=s; m_exps=e;
    }

    public IRStatement getStatement() {
	return m_stm;
    }

    public IRExpressionList getExpressionList() {
	return m_exps;
    }

    private IRStatement m_stm;
    private IRExpressionList m_exps;
}

public class Canonicalizer {

	/** The unique identifier of the function currently being canonicalized. */
	private int m_counterIndex;
    
	/** Path count of the canonicalized function. */
	private int m_pathCount;
	
	private String m_sourceFileName;

	private HashMap m_pathGraphs;

	private TraceAnalysis m_traceAnalysis;
	
	public Canonicalizer() {

    }

    public IRStatementList canonicalize(IRStatement prog, Frame frame) {

	BasicBlocker blocks = new BasicBlocker();
	TraceScheduler tracer = new TraceScheduler();

        m_curframe = frame;
	IRStatementList lineared = linearize(prog);

	blocks.formBlocks(lineared);

	/**HERE INSERT OPTIMIZATIONS**/
	
	// added for trace instrumentation and optimization
	new DominatorTree(blocks.getBlocks()).eliminateUnreachableBlocks();
	//end
	
	for(int i = 0; i < blocks.getBlocks().size(); i++) {
		System.out.println(blocks.getBlocks().get(i));
	}

	/* added for trace instrumentation */
	Trace hotPath;
	boolean complete = false;  // indicates whether the hot path that was found is complete
	PathGraph pathGraph;       // the path graph we're working on
	
	// always try to load any existing graph 
	pathGraph = getStoredPathGraph(blocks.getBlocks());
	// null indicates no stored graph
	boolean newGraph = pathGraph == null;
	if(newGraph) {
		// this is the first run
		System.out.println("Build new path graph for fct " + m_counterIndex);
		// this will create a partial or full path graph, depending on doIncrementalTracing()
		pathGraph = createInitialGraph(blocks.getBlocks());
		
		// if only one path then that's gotta be the hot path, otherwise we have no clue
		// (the initial graph contains at least two paths if branches or loops exist)
		complete = pathGraph.getPathCount() == 1;
		if(complete) {
			hotPath = pathGraph.getPath(0);
		}
		else {
			hotPath = new Trace();
		}
	}
	else {
		// this is a "next" run
		// with full instrumentation the hot path is found in one run
		// with incremental instrumentation the hot path can be incomplete
		int pathToExtend = computeHotPathFromTrace();
		if(pathToExtend < 0) {
			// function wasn't instrumented--hot path already known
			// TODO somehow find out what the hot path was and initialize hotPath accordingly
			complete = true;
			hotPath = new Trace();
		}
		else {
			hotPath = pathGraph.getPath(pathToExtend);
			complete = pathGraph.isComplete(pathToExtend);
			
			if(doFullTracing() && ! complete) {
				// just a sanity check for the full instrumentation scheme
				System.out.println("Full instrumentation didn't find complete hot path!");
			}
			
			// extension actually necessary
			if(complete == false) {
				System.out.println("Extend path " + pathToExtend + " for fct " + m_counterIndex);
				extendGraph(pathGraph, pathToExtend);
			}
		}
	}
	System.out.println(pathGraph.toString());
	// store the new graph immediately
	putPathGraph(pathGraph);
	// write out hot path information
	writeHotPath(hotPath, complete, newGraph == false);
	
	// compute instrumentation
	TraceInstrumentation instrument = new TraceInstrumentation(pathGraph, frame, m_counterIndex);
	m_pathCount = instrument.getPathCount();
	// instrument only if hot path not found yet
	if(complete == false) {
		instrument.placeInstrumentation();
	}		
	// if requested insert prelude
	if(m_prelude != null) {
		Block entry = instrument.uniqueEntry();
		for(int i = 0; i < m_prelude.getSize(); i++)
			entry.getBody().insertStatement(i, m_prelude.getStatement(i));
	}
	// if requested insert postlude
	if(m_postlude != null) {
		Block exit = instrument.uniqueExit();
		for(int i = 0; i < m_postlude.getSize(); i++)
			exit.getBody().insertStatement(exit.getBody().getSize() - 1, m_postlude.getStatement(i));
	}
	// trace instrumentation end
	
	for(int i = 0; i < blocks.getBlocks().size(); i++) {
		System.out.println(blocks.getBlocks().get(i));
	}

	/**END OPTIMIZATIONS**/

	tracer.formTraces(blocks);

	IRStatementList flattened = tracer.flattenTraces(tracer.getTraces()).append(blocks.getDoneLabel());
	
	return flattened;
    }
    
	/**
	 * Determines which path is the hot path given a trace file.
	 * @return The hot path according to the trace file.
	 */
	private int computeHotPathFromTrace() {
		// read trace file to find out what the most taken path was and return its number
		if(m_traceAnalysis == null)
			m_traceAnalysis = new TraceAnalysis();
		return m_traceAnalysis.getMostFrequentPath(m_counterIndex);
	}

	/**
	 * Writes information about the hot path to a file.
	 * @param hotPath
	 * @param complete
	 * @param appendToLastRun
	 */
	private void writeHotPath(Trace hotPath, boolean complete, boolean appendToLastRun) {
		try {
			String nameBase = m_sourceFileName + ".hot";
			if(doFullTracing() == false) nameBase += "_inc";
			
			PrintWriter pw;
			// write temp file with paths just of the current compiler run
			// hack: create new file if first function; append file if other functions
			pw = new PrintWriter(new FileWriter(nameBase + ".temp", m_counterIndex > 0));
			if(m_counterIndex == 0)
				pw.println("Paths");
			pw.println(m_counterIndex + " " + (complete ? "complete" : "partial") + " " + hotPath.toOneLine());
			pw.close();
			
			// write (append) results file with same content 
			// never append to last compiler run if full instrumentation
			if(doFullTracing()) appendToLastRun = false;
			// append the first function if the flag is set; otherwise always append
			pw = new PrintWriter(new FileWriter(nameBase, m_counterIndex > 0 || appendToLastRun));
			if(m_counterIndex == 0)
				pw.println("Paths");
			pw.println(m_counterIndex + " " + (complete ? "complete" : "partial") + " " + hotPath.toOneLine());
			pw.close();
		} 
		catch (IOException e) {
			System.err.println("Error writing out hot path.");
			e.printStackTrace();
		}
	}

	/**
	 * Extends a path in a graph with the next branch (resulting in more paths).
	 * @param incrementalGraph The graph to be modified.
	 * @param pathToExtend The path to be extended (value >= 0).
	 */
	private void extendGraph(PathGraph incrementalGraph, int pathToExtend) {
		incrementalGraph.extendPath(pathToExtend);
	}

	/**
	 * Creates an initial graph that used for the initial instrumentation.
	 * This graph should contain maximally one branch.
	 * @param blocks The blocks for the function.
	 * @return The initial path graph and never <code>null</code>.
	 */
	private PathGraph createInitialGraph(Vector blocks) {
		return new PathGraph(blocks, doFullTracing() == false);
	}

	/**
	 * This method indicates whether tracing is done incrementally or in one shot.
	 * @return <code>True</code> if tracing should happen incrementally, <code>false</code> otherwise.
	 */
	protected boolean doFullTracing() {
		// implementation analogous to Boolean.getBoolean()
		boolean result = false;
		try {
			result = System.getProperty("trace").equals("full");
		}
		catch(NullPointerException npe) {}
		catch(IllegalArgumentException iae) {}
		return result;
	}

	/**
	 * Stores a given path graph so that it can later be found as the
	 * graph stored for this function.
	 * @param incrementalGraph
	 */
	private void putPathGraph(PathGraph incrementalGraph) {
		if(m_pathGraphs == null) m_pathGraphs = new HashMap();
		// TODO do we need to prepare the graph for serialization?
		m_pathGraphs.put(new Integer(m_counterIndex), incrementalGraph);
	}
	
	public void storePathGraphs() {
		try {
			// any changes to the file format must be reflected in the read code in setSourceFileName
			ObjectOutputStream oos = new ObjectOutputStream(new FileOutputStream(getPathGraphFileName()));
			// "mark" path graph with source file name and instrumentation mode (full vs. incremental)
			oos.writeObject(m_sourceFileName);
			oos.writeBoolean(doFullTracing());
			// write all the path graphs
			oos.writeObject(m_pathGraphs);
			oos.close();
		} 
		catch (FileNotFoundException e) {
			System.err.println("Couldn't create file for storing path graphs.");
			e.printStackTrace();
		} 
		catch (IOException e) {
			System.err.println("Could not store path graphs.");
			e.printStackTrace();
		}
	}

	/**
	 * Returns any stored path graph (from a former compiler run) for the current function.
	 * @param blocks
	 * @return The stored path graph for the current function or <code>null</code> if none found.
	 */
	private PathGraph getStoredPathGraph(Vector blocks) {
		if(m_pathGraphs == null) return null;
		if(m_pathGraphs.containsKey(new Integer(m_counterIndex)) == false) return null;
		
		// find the right graph in the hash map 
		PathGraph result = (PathGraph) m_pathGraphs.get(new Integer(m_counterIndex));
		// TODO do we need to set the blocks or anything?
		result.setBlocks(blocks);
		return result;
	}

	public void setCounterIndex(int counterIndex) {
    	m_counterIndex = counterIndex;
    }
    
    public int getPathCount() {
    	return m_pathCount;
    }
    
    public void setSourceFileName(String fileName) {
    	this.m_sourceFileName = fileName;
    	m_pathGraphs = null;
    	try {
			ObjectInputStream ois = new ObjectInputStream(new FileInputStream(getPathGraphFileName()));
			if(ois.readObject().equals(fileName) && ois.readBoolean() == doFullTracing()) {
				m_pathGraphs = (HashMap) ois.readObject();
			}
			// else path graph file not suitable for this run--forces construction of new initial graph
			ois.close();
			return;
		} 
    	catch (FileNotFoundException e) {
    		System.out.println("No file with path graphs found.");
			//e.printStackTrace();
		} 
    	catch (IOException e) {
    		System.err.println("Problems reading path graphs file.");
			e.printStackTrace();
		} 
    	catch (ClassNotFoundException e) {
    		System.err.println("Missing class file for deserialization.");
			e.printStackTrace();
		}
    }

    /**
	 * @return
	 */
	private String getPathGraphFileName() {
		return "trace.pg";
	}

	private boolean isNop(IRStatement a) {
	return a.getClassName().equals("IRStatementExpression")
		&& ((IRStatementExpression)a).getExpression().getClassName().equals("IRConst");
    }

    private IRStatement seq(IRStatement a, IRStatement b) {

	if (isNop(a)) {
	    return b;
	}
	else if (isNop(b)) {
	    return a;
	}
	else {
	    return new IRSeq(a,b);
	}
    }

    private boolean commute(IRStatement a, IRExpression b) {
	return isNop(a) || b.getClassName().equals("IRConst");
    }

    private IRStatement do_stm(IRSeq s) {
	return seq(do_stm(s.getLeft()), do_stm(s.getRight()));
    }

    private IRStatement do_stm(IRMove s) {
	if (s.getDestination() instanceof IRTemp && s.getSource() instanceof IRCall) {
	    return reorder_stm(new MoveCall((IRTemp)s.getDestination(),(IRCall)s.getSource()));
	}
	if (s.getDestination().getClassName().equals("IRESeq")) {
	    return do_stm(new IRSeq(((IRESeq)s.getDestination()).getEffect(),new IRMove(((IRESeq)s.getDestination()).getExpression(),
		          s.getSource())));
	}
	else {
	    return reorder_stm(s);
	}
    }

    private IRStatement do_stm(IRStatementExpression s) {
	if (s.getExpression() instanceof IRCall) {
	    return reorder_stm(new ExpCall((IRCall)s.getExpression()));
	}
	else {
            return reorder_stm(s);
	}
    }

    private IRStatement do_stm(IRStatement s) {

	if (s.getClassName().equals("IRSeq")) {
	    return do_stm((IRSeq)s);
	}
	else if (s.getClassName().equals("IRMove")) {
	    return do_stm((IRMove)s);
	}
	else if (s.getClassName().equals("IRStatementExpression")) {
	    return do_stm((IRStatementExpression)s);
	}
	else {
	    return reorder_stm(s);
	}
    }

    private IRStatement reorder_stm(IRStatement s) {
	StatementExpressionList x = reorder(s.getComponents());
	return seq(x.getStatement(), s.build(x.getExpressionList()));
    }

    private IRESeq do_exp(IRESeq e) {

	IRStatement stms = do_stm(e.getEffect());
	IRESeq b = do_exp(e.getExpression());

	return new IRESeq(seq(stms,b.getEffect()), b.getExpression());
    }

    private IRESeq do_exp(IRExpression e) {
	if (e.getClassName().equals("IRESeq")) {
	    return do_exp((IRESeq)e);
	}
	else {
	    return reorder_exp(e);
	}
    }

    private IRESeq reorder_exp(IRExpression e) {
	StatementExpressionList x = reorder(e.getComponents());
	return new IRESeq(x.getStatement(), e.build(x.getExpressionList()));
    }

    private StatementExpressionList reorder(IRExpressionList exps) {

	if (exps.getSize() == 0) {
	    return nopNull;
	}
        else {
            IRExpression a = exps.getExpression(0);
            if (a instanceof IRCall) {
                Temp t = m_curframe.allocLocalTemp();
                IRExpression e = new IRESeq(new IRMove(new IRTemp(t), a),new IRTemp(t));
                return reorder(new IRExpressionList(e,exps.removeFront()));
            }
            else {
                IRESeq aa = do_exp(a);
                StatementExpressionList bb = reorder(exps.removeFront());
                if (commute(bb.getStatement(), aa.getExpression())) {
                    return new StatementExpressionList(seq(aa.getEffect(),bb.getStatement()),
                            new IRExpressionList(aa.getExpression(),bb.getExpressionList()));
                }
                else {
                    Temp t = m_curframe.allocLocalTemp();
                    return new StatementExpressionList(seq(aa.getEffect(),seq(new IRMove(new IRTemp(t),aa.getExpression()),bb.getStatement())),
                            new IRExpressionList(new IRTemp(t), bb.getExpressionList()));
                }
            }
	}
    }

    private IRStatementList linear(IRSeq s, IRStatementList l) {
	return linear(s.getLeft(),linear(s.getRight(),l));
    }
    private IRStatementList linear(IRStatement s, IRStatementList l) {

	if (s.getClassName().equals("IRSeq")) {
	    return linear((IRSeq)s, l);
	}
	else {
	    return new IRStatementList(s,l);
	}
    }

    private IRStatementList linearize(IRStatement s) {
	return linear(do_stm(s), null);
    }

    private Frame m_curframe;
    private StatementExpressionList nopNull = new StatementExpressionList(new IRStatementExpression(new IRConst(0)),null);
	private IRStatementList m_prelude;
	private IRStatementList m_postlude;

	/**
	 * @param prelude
	 */
	public void setPrelude(IRStatementList prelude) {
		m_prelude = prelude;
	}

	/**
	 * @param postlude
	 */
	public void setPostlude(IRStatementList postlude) {
		m_postlude = postlude;
	}
}
