use "test/alloctest_testGraphs.sml";
use "test/a_invariants.sml";

structure Set = Color.Set
structure Map = Color.Map

val registers = ref ["%eax","%ebx","%edx"]
(* Frame.registers *)

exception IllegalState

structure ListKey: ORD_KEY =
    struct
        type ord_key = int list
        fun compare(l1,l2) = 
            List.collate (fn (n1,n2) =>
                             if n1 < n2 then LESS
                             else if n1 > n2 then GREATER
                             else EQUAL) (l1,l2)
    end
structure ListSet = ListSetFn(ListKey)

fun genDummyGraph {numOfNodes,interferencePairList,movesPairList} graphName = 
    let fun tnodeBase tgPairList (temp:Temp.temp) = 
            let val igNodeOpt = 
                    List.find (fn (pTemp,_) => pTemp = temp) tgPairList
            in case igNodeOpt
                of SOME (_,igNode) => igNode
                 | _ => raise IllegalState
            end
        fun gtempBase tgPairList igNode = 
            let val tempOpt = 
                    List.find (fn (_,pIgNode) => Graph.eq(pIgNode,igNode)) 
                              tgPairList
            in case tempOpt
                of SOME (temp,_) => temp
                 | _ => raise IllegalState
            end
        fun makeIGraph numOfIgNodes =
            let val iGraph = Graph.newGraph()
                val _ = Temp.reset()
                val tgPairList = 
                    (List.tabulate (numOfIgNodes,
                                    fn n => (Temp.newtemp(), 
                                             Graph.newNode(iGraph))))
                    @ [(IA32Frame.EAX,Graph.newNode(iGraph)),
                       (IA32Frame.EDX,Graph.newNode(iGraph))]
            in (iGraph,tgPairList)
            end

        fun addInterferences igNodeList nPairList =
            let fun addInterferenceBase (n1,n2) =
                    Graph.mk_edge {from=List.nth(igNodeList,n1),
                                   to=  List.nth(igNodeList,n2)}
            in app addInterferenceBase nPairList
            end
        fun makeFGraphWithMoves igNodeList gtemp nPairList = 
            let val fGraphInitial = Flow.FGRAPH{control=Graph.newGraph(),
                                                def=Graph.Table.empty,
                                                use=Graph.Table.empty,
                                                ismove=Graph.Table.empty}
                fun addMove ((n1,n2),Flow.FGRAPH{control,ismove,def,use}) =
                    let val enter = Graph.Table.enter
                        val fgNode = Graph.newNode control
                        val (temp1,temp2) = (gtemp(List.nth(igNodeList,n1)),
                                             gtemp(List.nth(igNodeList,n2)))
                    in Flow.FGRAPH{control=control,
                                   def=enter(def,fgNode,[temp1]),
                                   use=enter(use,fgNode,[temp2]),
                                   ismove=enter(ismove,fgNode,true)}
                    end
            in foldl addMove fGraphInitial nPairList
            end
        val (iGraph,tgPairList) = makeIGraph numOfNodes
        val (_,igNodeList) = ListPair.unzip tgPairList
        val _ = addInterferences igNodeList interferencePairList
        val tnode = tnodeBase tgPairList
        val gtemp = gtempBase tgPairList
        val fGraph as Flow.FGRAPH{control,def,use,...}
          = makeFGraphWithMoves igNodeList gtemp movesPairList
        val fgNodes = Graph.nodes(control)
        fun getMoves fgNode = 
            let val defINode = tnode(hd(valOf(Graph.Table.look(def,fgNode))))
                val useINode = tnode(hd(valOf(Graph.Table.look(use,fgNode))))
            in {fgNode=fgNode,duPair=(defINode,useINode)}
            end
        val moves = List.map getMoves fgNodes
        val interference = Liveness.IGRAPH{graph=iGraph,
                                           tnode=tnode,
                                           gtemp=gtemp,
                                           moves=moves}
    in {fGraph=fGraph,fgNodes=fgNodes,
        interference=interference,graphName=graphName}
    end

fun igNodeAttrToStr(attr) = 
    case attr of
        Color.PRECOLORED =>   "Precolored"
      | Color.INITIAL =>      "Initial"
      | Color.SIMPLIFYWORK => "Simplifywork"
      | Color.FREEZEWORK =>   "Freezework"
      | Color.SPILLWORK =>    "Spillwork"
      | Color.SPILLED =>      "Spilled"
      | Color.COALESCED =>    "Coalesced"
      | Color.COLORED =>      "Colored"
      | Color.SELECTED =>     "Selected"

fun moveAttrToStr(attr) =
    case attr of
        Color.COALESCED_MOVE => "Coalesced"
      | Color.CONSTRAINED_MOVE => "Constrained"
      | Color.FROZEN_MOVE => "Frozen"
      | Color.WORKLIST_MOVE => "Worklist"
      | Color.ACTIVE_MOVE => "Active"

fun getIgNodeAttrBase igNodeAttrTbl node = 
    valOf(Graph.Table.look(igNodeAttrTbl, node))

fun getIgNodeStrWithAliasesBase gtemp getIgNodeAttr aliasNodeList node = 
    let val igNodeAttr = getIgNodeAttr node
        val aliasStr = 
            foldl (fn (n,s) => s ^ "(" ^ Main.makestring(gtemp(n)) ^ ")\\n")
                  "" aliasNodeList
    in "\"" ^ Main.makestring(gtemp(node)) ^ "\\n" ^ aliasStr ^
       igNodeAttrToStr(igNodeAttr) ^ "\""
    end

fun getIgNodeStrBase gtemp getIgNodeAttr node = 
    getIgNodeStrWithAliasesBase gtemp getIgNodeAttr nil node

fun getIgNodeExtAttrStrBase igNodeAttrTbl igNode = 
    let val igNodeAttr = getIgNodeAttrBase igNodeAttrTbl igNode
    in case igNodeAttr 
        of Color.PRECOLORED => "color=black,fontcolor=black"
         | Color.INITIAL => "color=black,fontcolor=black"
         | Color.SELECTED => "color=lightgray,fontcolor=lightgray"
         | Color.COALESCED => "color=lightgray,fontcolor=lightgray"
         | _ => "color=black,fontcolor=black"
    end

fun showIgNode (cEnv as Color.ENV{igNodeList,tnode,gtemp,...})
               (aliasTbl,igNodeAttrTbl) = 
    let val getIgNodeAttr= getIgNodeAttrBase igNodeAttrTbl
        val getIgNodeStr = getIgNodeStrBase gtemp getIgNodeAttr
        val getIgNodeExtAttrStr = getIgNodeExtAttrStrBase igNodeAttrTbl
        fun showIgNode0 igNode = 
            let val aliasNodeList = 
                    Color.getAliasNodeList(aliasTbl,igNodeList,igNode)
                val getIgNodeStrLabel = getIgNodeStrWithAliasesBase
                                            gtemp getIgNodeAttr aliasNodeList
            in "\t" ^ getIgNodeStr igNode ^ 
               "[label=" ^ getIgNodeStrLabel igNode ^ "," ^
               getIgNodeExtAttrStr igNode ^ "]\n"
            end
    in concat(map showIgNode0 igNodeList)
    end

fun showAdjIgNodes (cEnv as Color.ENV{igNodeList,tnode,gtemp,...})
                   igNodeAttrTbl =
    let fun getTempPairSet igNode =
            let val getAdjacentTemps = 
                    Color.getAdjacentTempsBase cEnv igNodeAttrTbl
                val t1 = gtemp igNode
                val adjTempsList = Set.listItems(getAdjacentTemps igNode)
                fun tempPair(temp) = 
                    Set.listItems(Set.add(Set.add(Set.empty,t1),temp))
                val tpList = map tempPair adjTempsList
            in foldl ListSet.add' ListSet.empty tpList
            end
        val tempPairSetList = map getTempPairSet igNodeList
        val tempPairList = 
            ListSet.listItems(foldl ListSet.union 
                                    ListSet.empty tempPairSetList)
        fun showAdjIgNodes0 [t1,t2] =
            let val n1 = tnode t1 (* XXX *)
                val n2 = tnode t2
                val getIgNodeAttr= getIgNodeAttrBase igNodeAttrTbl
                val getIgNodeStr = getIgNodeStrBase gtemp getIgNodeAttr
                fun isSorC(igNodeAttr) = igNodeAttr = Color.SELECTED orelse 
                                         igNodeAttr = Color.COALESCED
                val n1Attr = getIgNodeAttr n1
                val n2Attr = getIgNodeAttr n2
            in if isSorC(n1Attr) orelse isSorC(n2Attr)
               then ""
               else "\t" ^ getIgNodeStr n1 ^  " -- " ^ getIgNodeStr n2 ^ "\n"
            end
          | showAdjIgNodes0 _ = raise IllegalState
    in concat(map showAdjIgNodes0 tempPairList)
    end

fun showMovesNodes (cEnv as Color.ENV{gtemp,moves,...})
                   (moveAttrTbl,igNodeAttrTbl) =
    let val getIgNodeAttr= getIgNodeAttrBase igNodeAttrTbl
        val getIgNodeStr = getIgNodeStrBase gtemp getIgNodeAttr
        fun showMovesNodes0 {fgNode,duPair=(defNode,useNode)} = 
            let val moveAttr = valOf(Graph.Table.look(moveAttrTbl,fgNode))
                val moveAttrStr = moveAttrToStr(moveAttr)
            in ("\t" ^ getIgNodeStr defNode ^ " -- " ^ getIgNodeStr useNode ^ 
                " [style=dashed,color=red,fontcolor=red,label=\"" ^ 
                moveAttrStr ^ "\"]\n")
        end
    in concat(map showMovesNodes0 moves)
    end

fun writeGraph out (cEnv,(aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack)) =
    let fun output(str) = TextIO.output(out,str)
        val igNodeStr = showIgNode cEnv (aliasTbl,igNodeAttrTbl)
        val adjIgNodesStr = showAdjIgNodes cEnv igNodeAttrTbl
        val movesStr = showMovesNodes cEnv (moveAttrTbl,igNodeAttrTbl)
    in (output("graph G {\n");
(*        output("\tsize=\"3,3\"\n"); *)
        output("\tnode[fontsize=10,width=.1,height=.05]\n");
        output("\tedge[fontsize=10,len=2]\n");
        output(igNodeStr);
        output(adjIgNodesStr);
        output(movesStr);
        output("}\n"))
    end

fun doInitial {fGraph,fgNodes,interference,graphName} =
    let val Flow.FGRAPH{control,ismove,def,use} = fGraph
        val Liveness.IGRAPH{graph,tnode,gtemp,moves} = interference 
        val moveAttrTbl = 
            let val initMoveAttr = RegAlloc.initMoveAttrBase ismove
                val fgNodeList = Graph.nodes(control)
            in foldl initMoveAttr Graph.Table.empty fgNodeList
            end
        val igNodeList = Graph.nodes(graph)
        val movesTbl = 
            let val makeMovesTable = 
                    RegAlloc.makeMovesTableBase ismove def use tnode
            in foldl makeMovesTable Graph.Table.empty fgNodes
            end
        val aliasTbl = Graph.Table.empty:(Graph.node Graph.Table.table)
        val cEnv = Color.ENV{movesTbl=movesTbl,moves=moves,
                             igNodeList=igNodeList,
                             gtemp=gtemp,tnode=tnode,registers=(!registers)}
        val igNodeAttrTblInitial =
            Color.getInitialIgNodeAttrTbl(gtemp,Frame.tempMap,igNodeList)
        val igNodeAttrTbl =
            Color.makeWorkList cEnv (igNodeAttrTblInitial,moveAttrTbl) 
        val selectStack = nil
        val tables = (aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack)
        fun writeToFile() = Main.withOpenFile (graphName ^ "_0.nto")
                                   (fn out => writeGraph out (cEnv,tables))
    in (print("-- doInitial --\n");
        writeToFile();
        invariantTest cEnv (igNodeAttrTbl,moveAttrTbl);
        (cEnv,tables))
    end

fun getWorkLists (cEnv as Color.ENV{igNodeList,moves,...}) 
                 (moveAttrTbl,igNodeAttrTbl) =
    let val getNodesWithAttr = 
            Color.getNodesWithAttrBase igNodeList igNodeAttrTbl 
        val simplifyWorkList = getNodesWithAttr Color.SIMPLIFYWORK
        val workListMoves = 
            let fun isWorkListMove {fgNode,duPair} =
                    valOf(Graph.Table.look(moveAttrTbl,fgNode)) =
                    Color.WORKLIST_MOVE
            in List.filter isWorkListMove moves
            end
        val freezeWorkList = getNodesWithAttr Color.FREEZEWORK
        val spillWorkList = getNodesWithAttr Color.SPILLWORK
    in (simplifyWorkList,workListMoves,freezeWorkList,spillWorkList)
    end

fun execLoop _ _ _ tables (nil,nil,nil,nil) = tables
  | execLoop graphName loopCount
             (cEnv as Color.ENV{igNodeList,moves,...}) tables 
             (simplifyWorkList,workListMoves,freezeWorkList,spillWorkList) = 
    let val tables as (aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack) =
            if null(simplifyWorkList) = false
            then Color.simplify cEnv tables simplifyWorkList
            else if null(workListMoves) = false
            then Color.coalesce cEnv tables workListMoves
            else if null(freezeWorkList) = false
            then Color.freeze cEnv tables freezeWorkList
            else if null(spillWorkList) = false
            then Color.selectSpill cEnv tables spillWorkList
            else raise IllegalState
        val workLists = getWorkLists cEnv (moveAttrTbl,igNodeAttrTbl)
        exception LoopCountExceeded
        fun writeToFile() = 
            if loopCount > 25 then raise LoopCountExceeded
            else Main.withOpenFile (graphName ^ "_" ^
                                    Int.toString(loopCount) ^ ".nto")
                                   (fn out => writeGraph out (cEnv,tables))
    in (writeToFile();
        execLoop graphName (loopCount+1) cEnv tables workLists)
    end

fun execTest testGraph graphName =
    let val graphs as {fGraph,fgNodes,interference,graphName} = 
            genDummyGraph testGraph graphName
        val (cEnv,tables as (_,moveAttrTbl,igNodeAttrTbl,_)) = doInitial graphs
        val workLists = getWorkLists cEnv (moveAttrTbl,igNodeAttrTbl)
        val (aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack) =
            execLoop graphName 1 cEnv tables workLists
    in print("-- execTest done --\n")
    end

(* A test just to make sure RegAlloc.alloc won't crash *)
fun regAllocTest filename =
    let val instrsAndFrameList = Main.getInstrsAndFrameList filename
        val instrsAndFrame = hd(instrsAndFrameList)
        val result = RegAlloc.alloc instrsAndFrame
    in (print "regAllocTest done\n") 
    (* discarding result for now - will be checked in future tests *)
    end