package edu.cmu.cs.l3.regalloc;

import edu.cmu.cs.l3.frame.*;
import edu.cmu.cs.l3.assem.*;
import edu.cmu.cs.l3.general.*;
import edu.cmu.cs.l3.graph.*;
import edu.cmu.cs.l3.flowgraph.*;
import java.util.*;
import edu.cmu.cs.l3.temp.*;
import edu.cmu.cs.l3.frame.*;

public abstract class RegisterAllocator implements TempMap {

    public class Stats {
        public Stats() { }
        public int m_spills=0, m_delete = 0, m_iterations=0, m_cnodes=0;
    }

    public static Stats m_stats=null;

    public RegisterAllocator(IGConstructor cons, FGConstructor fgcons, RegisterAllocatorOptions options) {
        m_igcons = cons; m_options = options; m_fgcons = fgcons;
        if (m_stats == null) m_stats = new Stats();
    }

    protected abstract boolean rewrite(Frame frame, List newcode, List code, List spills);

    public List allocate(Frame frame, List icode, List registers) throws GraphException {
        int i;
        List acode;

        acode = regalloc(frame,icode,registers);
        // Kill redundant moves (coalescing)
        for (i = 0; i < acode.size(); i++) {
            Instruction instr = (Instruction) acode.get(i);
            if (instr instanceof MoveInstruction) {
                MoveInstruction move = (MoveInstruction) instr;
                if (!move.isMemRelated() && move.sourceTemps().size() > 0 &&
                    move.destTemps().size() > 0) {
                    Temp stemp = (Temp)move.sourceTemps().get(0);
                    Temp dtemp = (Temp)move.destTemps().get(0);
                    if (stemp.isSpilled() || stemp.isGlobal() || dtemp.isSpilled() ||
                        dtemp.isGlobal() || stemp.hasValue() || dtemp.hasValue()) continue;
                    else if (((String)m_tempmap.get(stemp)).equals(m_tempmap.get(dtemp))) {
                        acode.remove(instr); i--; m_stats.m_delete++;
                    }
                }
            }
        }
        m_stats.m_cnodes = m_coalescedNodes.size();
        return acode;
    }

    private List regalloc(Frame frame, List code, List registers) throws GraphException {

        List altercode = new ArrayList();

        m_stats.m_iterations++;

        m_tempmap = new Hashtable();

        FlowGraph fg = m_fgcons.build(code);
        build(fg,registers,frame.getFrameTemp());

        m_select = new Stack(); m_initial = new ArrayList();
        m_initial.addAll(m_ig.getVertices());

        buildWorklists(registers.size());
        while (m_simplify.size() > 0 || m_spill.size() > 0
               || (m_options.m_coalesce && m_freeze.size() > 0)
               || (m_options.m_coalesce && m_moves.size() > 0)) {
            if (m_simplify.size() > 0)
                simplify(registers.size());
            else if (m_options.m_coalesce && m_moves.size() > 0)
                coalesce(frame, registers.size());
            else if (m_options.m_coalesce && m_freeze.size() > 0)
                freeze(registers.size());
            else if (m_spill.size() > 0)
                selectspill(registers.size());
        }

        List spills = color(frame,registers);
        m_stats.m_spills += spills.size();
        if (spills.size() > 0) {
            if (!rewrite(frame, altercode, code, spills)) {
                return altercode;
            }
            return regalloc(frame, altercode, registers);
        }
        else
            return code;
    }

    private List color(Frame frame, List acolors) throws GraphException {
        List okcolors = new ArrayList();
        List spills = new ArrayList();
        int i;

        List verts = m_ig.getVertices();
        for (i = 0; i < verts.size(); i++) {
            InterferenceVertex v = (InterferenceVertex) verts.get(i);
            List adjs = v.getAdjacents();
            if (v.isPrecolored())
                m_tempmap.put(v.getTemp(),frame.tempMap(v.getColor()));
        }

        while (!m_select.empty()) {
            okcolors = new ArrayList();
            okcolors.addAll(acolors);
            InterferenceVertex v = (InterferenceVertex) m_select.pop();

            m_ig.reviveVertex(v);
            List adjs = v.getAdjacents();
            for (i = 0; i < adjs.size(); i++) {
                InterferenceVertex u = (InterferenceVertex) adjs.get(i);
                u = (InterferenceVertex) getalias(u);
                if (u.isColored() || u.isPrecolored()) {
                    okcolors.remove(u.getColor());
                }
            }

            if (v.isPrecolored()) {
                m_tempmap.put(v.getTemp(),frame.tempMap(v.getColor()));
            }
            else if (okcolors.size() > 0) {
                Random rnd = new Random();
                int c = rnd.nextInt(okcolors.size());
                v.setColor((Temp)okcolors.get(c));
                m_tempmap.put(v.getTemp(),frame.tempMap((Temp)okcolors.get(c)));
                okcolors.remove(c);
            }
            else {
                spills.add(v);
            }
        }

        for (i = 0; i < m_coalescedNodes.size(); i++) {
            InterferenceVertex v = (InterferenceVertex) m_coalescedNodes.get(i);
            InterferenceVertex av = (InterferenceVertex) getalias(v);

            Temp c = av.getColor();
            if (c != null) {
                v.setColor(c);
                m_tempmap.put(v.getTemp(), frame.tempMap(c));
            }
            else
                spills.add(v);
        }

        return spills;
    }

    private void freeze(int k) throws GraphException {
        Vertex v = (Vertex) m_freeze.get(0);
        m_freeze.remove(0);
        m_simplify.add(v);
        freezemoves(v,k);
    }

    private void freezemoves(Vertex u, int k) throws GraphException {
        List moves = ((InterferenceVertex)u).getMoveList();
        int i;
        Vertex x,y,v;

        for (i = 0; i < moves.size(); i++) {
            MoveInstruction m = (MoveInstruction) moves.get(i);
            Temp tx = (Temp) m.sourceTemps().get(0);
            Temp ty = (Temp) m.destTemps().get(0);

            x = m_ig.getIntVertex(tx.num(), new ArrayList());
            y = m_ig.getIntVertex(ty.num(), new ArrayList());
            x = getalias(x); y = getalias(y); u = getalias(u);
            if (y == u)
                v = x;
            else
                v = y;

            m_activeMoves.remove(m);
            if (((InterferenceVertex)v).getMoveList().size() == 0 &&
                v.getDegree(Graph.DEGREE) < k) {
                m_freeze.remove(v);
                m_simplify.add(v);
            }
        }
    }

    private void simplify(int k) throws GraphException {
        List adjs;
        int i;

        Vertex v = (Vertex) m_simplify.get(0);
        m_simplify.remove(0);
        m_select.push(v);

        adjs = m_ig.getAdjacentsWithSuspend(v);
        m_ig.suspendVertex(v);

        for (i = 0; i < adjs.size(); i++)
            testDegree((InterferenceVertex)adjs.get(i), k);
    }

    private void combine(InterferenceVertex u, InterferenceVertex v, int k)
        throws GraphException {
        int i;

        if (m_freeze.contains(v))
            m_freeze.remove(v);
        else
            m_spill.remove(v);

        m_coalescedNodes.add(v);
        m_alias[v.getID()] = u;
        u.getMoveList().addAll(v.getMoveList());

        List adjs = m_ig.getAdjacentsWithSuspend(v);
        m_ig.suspendVertex(v);

        for (i = 0; i < adjs.size(); i++) {
            InterferenceVertex t = (InterferenceVertex) adjs.get(i);
            m_ig.createIntEdge(t,u);
            testDegree(t,k);
        }
        if (!u.isPrecolored() && u.getDegree(Graph.DEGREE) >= k && m_freeze.contains(u)) {
            m_freeze.remove(u);
            m_spill.add(u);
        }
    }

    private void coalesce(Frame frame, int k) throws GraphException {
        MoveInstruction m;
        InterferenceVertex x,y,u,v;
        int i;

        m = (MoveInstruction) m_moves.get(0); m_moves.remove(0);
        Temp tx = (Temp) m.sourceTemps().get(0);
        Temp ty = (Temp) m.destTemps().get(0);

        if (tx == frame.getFrameTemp() || ty == frame.getFrameTemp() ||
            tx.isSpilled() || ty.isSpilled() || tx.hasValue()
            || ty.hasValue() || m.isMemRelated()) return;

        x = m_ig.getIntVertex(tx.num(), new ArrayList());
        y = m_ig.getIntVertex(ty.num(), new ArrayList());
        x = (InterferenceVertex) getalias(x);
        y = (InterferenceVertex) getalias(y);

        if (y.isPrecolored()) {
            u = y; v = x;
        }
        else {
            u = x; v = y;
        }

        if (u == v) {
            testForSimplifyAdd(u, k);
        }
        else if (v.isPrecolored() || u.isAdjacent(v)) {
            testForSimplifyAdd(u, k);
            testForSimplifyAdd(v, k);
        }
        else {
            if (u.isPrecolored()) {
                List vadj = m_ig.getAdjacentsWithSuspend(v);
                for (i = 0; i < vadj.size(); i++) {
                    Vertex z = (Vertex) vadj.get(i);
                    if (!precoloredTest(z, u, k)) {
                        m_activeMoves.add(m); return;
                    }
                }

            //    System.out.println(m_coalescedNodes);
            //    System.out.println("\t"+u + "(" + u.getID() + ") " + v + "(" + v.getID() + ")");
             //   System.out.print(m.format(new DefaultMap()));
                combine((InterferenceVertex)u, (InterferenceVertex)v, k);
                testForSimplifyAdd(u, k);
            }
            else {
                List adjs = m_ig.getAdjacentsWithSuspend(u);
                adjs.addAll(m_ig.getAdjacentsWithSuspend(v));

                if (conservativeTest(adjs, k)) {
               //     System.out.println(m_coalescedNodes);
               ///     System.out.println("\t"+u + "(" + u.getID() + ")" + v + "(" + v.getID() + ")");
               //     System.out.print(m.format(new DefaultMap()));
                    combine((InterferenceVertex)u,(InterferenceVertex)v,k);
                    testForSimplifyAdd(u, k);
                }
                else {
                    m_activeMoves.add(m); return;
                }
            }
        }
    }

    private boolean precoloredTest(Vertex v, Vertex u, int k)
        throws GraphException {
        return (v.getDegree(Graph.DEGREE) < k || ((InterferenceVertex)v).isPrecolored()
                || v.isAdjacent(u));
    }

    private boolean conservativeTest(List vertices, int k) throws GraphException {
        int i,crit=0;
        for (i = 0; i < vertices.size(); i++) {
            Vertex v = (Vertex) vertices.get(i);
            if (v.getDegree(Graph.DEGREE) >= k) crit++;
        }
        return (crit < k);
    }

    private void testForSimplifyAdd(Vertex v, int k) throws GraphException {
        if (!((InterferenceVertex)v).isPrecolored() &&
            (m_options.m_coalesce && !isMoveRelated((InterferenceVertex)v)) &&
            v.getDegree(Graph.DEGREE) < k) {
            m_freeze.remove(v);
            m_simplify.add(v);
        }
    }

    private Vertex getalias(Vertex v) {
        if (m_coalescedNodes.contains(v)) {
            return getalias(m_alias[v.getID()]);
        }
        else
            return v;
    }

    private void selectspill(int k) throws GraphException {
        int i,bspill=-1;
        double bvalue=0;

        for (i = 0; i < m_spill.size(); i++) {
            InterferenceVertex v = (InterferenceVertex) m_spill.get(i);
            if (v.getTemp().isMemHelper()) continue;
            double dvalue = (double) v.getDegree(Graph.DEGREE);
            double udvalue = (double) v.getTemp().getNumDefs() + v.getTemp().getNumUses();
            if (bspill < 0) {
                bvalue = (2.0*udvalue)/dvalue;
                bspill = i;
                continue;
            }
            else {
                if ((udvalue/dvalue)<bvalue) {
                    bvalue = udvalue/dvalue;
                    bspill = i;
                }
            }
        }

        if (bspill < 0) bspill = 0;

        Vertex m = (Vertex) m_spill.get(bspill);
        m_simplify.add(m);
        m_spill.remove(bspill);

        if (m_options.m_coalesce) freezemoves(m, k);
    }

    private void testDegree(InterferenceVertex v, int k) throws GraphException {
        List enables;
        if (v.getDegree(Graph.DEGREE) == k && !v.isPrecolored()) {
            if (m_options.m_coalesce) {
                enables = m_ig.getAdjacentsWithSuspend(v);
                enables.add(v);
                enableMoves(enables);
            }
            m_spill.remove(v);
            if (m_options.m_coalesce && isMoveRelated((InterferenceVertex)v))
                m_freeze.add(v);
            else
                m_simplify.add(v);
        }
    }

    private void enableMoves(List verts) throws GraphException {
        int i,j;
        for (i = 0; i < verts.size(); i++) {
            Vertex v = (Vertex) verts.get(i);
            List moves = getVertexMoves((InterferenceVertex)v);
            for (j = 0; j < moves.size(); j++) {
                Instruction instr = (Instruction) moves.get(j);
                if (m_activeMoves.contains(instr)) {
                    m_activeMoves.remove(instr);
                    m_moves.add(instr);
                }
            }
        }
    }

    private List getVertexMoves(InterferenceVertex v) throws GraphException {
        List moves = new ArrayList(),actmoves = new ArrayList();
        List res = new ArrayList();
        int i,j;

        moves = v.getMoveList();
        actmoves.addAll(m_activeMoves); actmoves.addAll(m_moves);

        for (i = 0; i < moves.size(); i++) {
            MoveInstruction move = (MoveInstruction) moves.get(i);
            if (actmoves.contains(move))
                res.add(move);
        }

        return res;
    }

    private boolean isMoveRelated(InterferenceVertex v) throws GraphException {
        return (getVertexMoves(v).size() != 0);
    }

    private void buildWorklists(int k) throws GraphException {
        int i;

        for (i = 0; i < m_initial.size(); i++) {
            InterferenceVertex v = (InterferenceVertex) m_initial.get(i);

            if (v.isPrecolored())
                continue;
         //       m_simplify.add(v);
            else if (v.getDegree(Graph.DEGREE) >= k) {
                m_spill.add(v);
            }
            else if (m_options.m_coalesce && isMoveRelated(v)) {
                m_freeze.add(v);
            }
            else {
                m_simplify.add(v);
            }
        }
        m_initial.clear();
    }

    private void build(FlowGraph code, List registers, Temp fp) throws GraphException {

        List precols = new ArrayList();
        precols.addAll(registers); precols.add(fp);

        m_ig = m_igcons.build(code,precols,m_moves);
        m_alias = new InterferenceVertex[m_ig.getNumVertices()];
        m_coalescedNodes = new ArrayList(); m_activeMoves = new ArrayList();
    }

    protected Dictionary m_tempmap = new Hashtable();
    protected RegisterAllocatorOptions m_options;

    private IGConstructor m_igcons;
    private FGConstructor m_fgcons;
    private InterferenceGraph m_ig;
    private Stack m_select = new Stack();
    private Vertex[] m_alias;
    private List m_simplify = new ArrayList(), m_spill = new ArrayList();
    private List m_initial = new ArrayList(), m_moves = new ArrayList();
    private List m_freeze = new ArrayList();
    private List m_activeMoves = new ArrayList();
    private List m_coalescedNodes = new ArrayList();
    private int m_its=0;

}