signature COLOR =
sig
  structure Frame : FRAME
  type allocation = Frame.register Temp.Table.table
  datatype moveAttr = COALESCED_MOVE | CONSTRAINED_MOVE | FROZEN_MOVE
                    | WORKLIST_MOVE | ACTIVE_MOVE
  datatype env = ENV of {igNodeList: Graph.node list,
                         movesTbl: (Graph.node list) Graph.Table.table,
                         moves: {fgNode: Graph.node,
                                 duPair: (Graph.node * Graph.node)} list,
                         gtemp: Graph.node -> Temp.temp,
                         tnode: Temp.temp -> Graph.node,
                         registers: Frame.register list}
  val color : {interference: Liveness.igraph,
               initial: allocation,
               spillCost: Graph.node -> int,
               registers: Frame.register list,
               moveAttrTbl: moveAttr Graph.Table.table,
               movesTbl: (Graph.node list) Graph.Table.table}
              -> allocation * Temp.temp list
end

structure Color (* :COLOR *)  =
struct
  exception IllegalState
  structure Frame = Frame
  structure Table = Graph.Table
  structure Set = IntListSet
  structure Map = IntBinaryMap
  type allocation = Frame.register Temp.Table.table
  datatype igNodeAttr = 
           PRECOLORED | INITIAL | SIMPLIFYWORK 
         | FREEZEWORK | SPILLWORK | SPILLED | COALESCED 
         | COLORED | SELECTED
  datatype moveAttr = COALESCED_MOVE | CONSTRAINED_MOVE | FROZEN_MOVE
                    | WORKLIST_MOVE | ACTIVE_MOVE
  datatype env = ENV of {igNodeList: Graph.node list,
                         movesTbl: (Graph.node list) Graph.Table.table,
                         moves: {fgNode: Graph.node,
                                 duPair: (Graph.node * Graph.node)} list,
                         gtemp: Graph.node -> Temp.temp,
                         tnode: Temp.temp -> Graph.node,
                         registers: Frame.register list}
  structure NodeKey: ORD_KEY =
    struct
      type ord_key = Graph.node
      fun compare((_,a),(_,b)) = if a < b then LESS
                                 else if a > b then GREATER
                                 else EQUAL
    end
  structure NodeSet = ListSetFn(NodeKey)

  fun getInitialIgNodeAttrTbl(gtemp,initial,igNodeList):(igNodeAttr Table.table) = 
      let fun initIgNodeAttrTbl(igNode,table) =
              let val isPrecolored = 
                      Option.isSome(Temp.Table.look(initial,gtemp(igNode)))
              in if isPrecolored = true 
                 then Table.enter(table,igNode,PRECOLORED)
                 else Table.enter(table,igNode,INITIAL)
              end
      in foldl initIgNodeAttrTbl Table.empty igNodeList
      end

  fun getAlias(aliasTbl,igNode) =
      case Table.look(aliasTbl,igNode) 
       of NONE => igNode
        | SOME aliasNode => getAlias(aliasTbl,aliasNode)

  fun getAliasNodeList(aliasTbl,igNodeList,igNode) =
      List.filter (fn n => Graph.eq(n,igNode) = false andalso
                           Graph.eq(getAlias(aliasTbl,n),igNode)) 
                  igNodeList

  fun getNodesWithAttrBase igNodeList igNodeAttrTbl (attr:igNodeAttr) =
      let fun hasAttr(igNode) =
              valOf(Graph.Table.look(igNodeAttrTbl,igNode)) = attr
              handle e => raise e
      in List.filter hasAttr igNodeList
      end 

  fun getAdjacentTempsBase (cEnv as ENV{igNodeList,gtemp,...})
                           igNodeAttrTbl n = 
      let fun getTempsWithAttr(attr) = 
              List.map gtemp (getNodesWithAttrBase igNodeList igNodeAttrTbl attr)
                       handle e => raise e
          fun makeNewSet(list) = Set.addList(Set.empty,list)
          val adjList = Graph.adj(n)
          val sTempSet = makeNewSet(getTempsWithAttr(SELECTED))
          val scTempSet = Set.addList(sTempSet,getTempsWithAttr(COALESCED))
          val adjTempSet = makeNewSet(List.map gtemp adjList)
      in Set.difference(adjTempSet,scTempSet)
      end

  fun getAdjacentNodesBase (cEnv as ENV{igNodeList,...}) igNodeAttrTbl n = 
      let val getNodesWithAttr = getNodesWithAttrBase igNodeList igNodeAttrTbl 
          fun makeNewSet(list) = NodeSet.addList(NodeSet.empty,list)
          val sNodeSet = makeNewSet(getNodesWithAttr(SELECTED))
          val scNodeSet = NodeSet.addList(sNodeSet,getNodesWithAttr(COALESCED))
          val adjNodeSet = makeNewSet(Graph.adj(n))
      in NodeSet.difference(adjNodeSet,scNodeSet)
      end

  fun degreeBase (cEnv as ENV{igNodeList,registers,...}) 
                 igNodeAttrTbl n = 
      if (valOf(Graph.Table.look(igNodeAttrTbl,n))) = PRECOLORED
      (* returns "virtually infinite" degree for precolored node *)
      then length(igNodeList) + length(registers) 
      else NodeSet.numItems(getAdjacentNodesBase cEnv igNodeAttrTbl n)

  fun nodeMoves igNodeList (aliasTbl,movesTbl,moveAttrTbl) n =
      let fun nodeMoves0 n =
              let val moveList = getOpt(Table.look(movesTbl,n),nil)
                  fun isActiveMovesOrWorkListMoves moveNode =
                      let val moveAttr = 
                              valOf(Table.look(moveAttrTbl,moveNode))
                              handle e => raise e
                      in moveAttr = ACTIVE_MOVE orelse
                         moveAttr = WORKLIST_MOVE
                      end
              in List.filter isActiveMovesOrWorkListMoves moveList 
              end
          val aliasList = getAliasNodeList(aliasTbl,igNodeList,n)
          val result = List.map nodeMoves0 (n::aliasList)
          (* caution - result may contain duplicate nodes *)
      in List.concat result
      end
  fun moveRelatedBase igNodeList (aliasTbl,movesTbl,moveAttrTbl) igNode = 
      length(nodeMoves igNodeList (aliasTbl,movesTbl,moveAttrTbl) igNode) <> 0

  fun makeWorkList (cEnv as ENV{igNodeList,movesTbl,registers,...})
                   (igNodeAttrTbl,moveAttrTbl) = 
      let val degree = degreeBase cEnv igNodeAttrTbl
          (* makeWorkList does not need aliasTbl for "moveRelated" - 
             using Table.empty for substitute here *)
          val moveRelated = 
              moveRelatedBase igNodeList (Table.empty,movesTbl,moveAttrTbl)

          fun setAttr(igNode,igNodeAttrTbl) = 
              let val attr = valOf(Table.look(igNodeAttrTbl,igNode))
                      handle e => raise e
                  val K = length(registers)
              in case attr
                  of INITIAL =>
                     if degree(igNode) >= K
                     then Table.enter(igNodeAttrTbl,igNode,SPILLWORK)
                     else if moveRelated(igNode)
                     then Table.enter(igNodeAttrTbl,igNode,FREEZEWORK)
                     else Table.enter(igNodeAttrTbl,igNode,SIMPLIFYWORK)
                   | _ => igNodeAttrTbl
              end
      in foldl setAttr igNodeAttrTbl igNodeList
      end

  fun enableMoves (movesTbl,moveAttrTbl) nodes =
      let fun enableMoves0(n,moveAttrTbl) =
              let val nodeMovesN = getOpt(Table.look(movesTbl,n),nil)
                  fun updateMoveAttrTbl(m,moveAttrTbl) =
                      let val attr = valOf(Table.look(moveAttrTbl,m))
                              handle e => raise e
                      in if attr = ACTIVE_MOVE
                         then Table.enter(moveAttrTbl,m,WORKLIST_MOVE)
                         else moveAttrTbl
                      end
              in foldl updateMoveAttrTbl moveAttrTbl nodeMovesN
              end
      in foldl enableMoves0 moveAttrTbl nodes
      end

  fun decrementDegreeBase (cEnv as ENV{movesTbl,registers,igNodeList,...}) 
                          aliasTbl
                          (m,(igNodeAttrTbl,moveAttrTbl)) =
                          (* the order of third set of arguments 
                             (m,(igNodeAttrTbl,moveAttrTbl))
                             is arranged for use with List.foldl *)
      let val degree = degreeBase cEnv igNodeAttrTbl
          val K = length(registers)
      in if degree(m) = K - 1
         then let val adjNodeList = 
                      NodeSet.listItems(getAdjacentNodesBase cEnv igNodeAttrTbl m)
                  val moveAttrTbl' = enableMoves (movesTbl,moveAttrTbl) (m::adjNodeList)
                  val moveRelated = 
                      moveRelatedBase igNodeList (aliasTbl,movesTbl,moveAttrTbl')
              in if moveRelated m = true 
                 then (Table.enter(igNodeAttrTbl,m,FREEZEWORK),moveAttrTbl') 
                 else (Table.enter(igNodeAttrTbl,m,SIMPLIFYWORK),moveAttrTbl')
              end
         else (igNodeAttrTbl,moveAttrTbl)
      end

  fun simplify cEnv (aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack) simplifyWorkList =
      let val n = hd(simplifyWorkList)
          val igNodeAttrTbl' = Table.enter(igNodeAttrTbl,n,SELECTED)
          val selectStack' = n::selectStack
          val moveAttrTbl' = moveAttrTbl
          val adjNodeList = NodeSet.listItems(getAdjacentNodesBase cEnv igNodeAttrTbl' n)
          val decrementDegree = decrementDegreeBase cEnv aliasTbl
          val (igNodeAttrTbl'',moveAttrTbl'') =
              foldl decrementDegree (igNodeAttrTbl',moveAttrTbl') adjNodeList
      in (aliasTbl,moveAttrTbl'',igNodeAttrTbl'',selectStack')
      end

  fun coalesce (cEnv as ENV{igNodeList,movesTbl,registers,...})
               (aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack) 
               workListMoves =
      let val K = length(registers)
          fun addWorkList(moveAttrTbl,igNodeAttrTbl,u) = 
              let val moveRelated = 
                      moveRelatedBase 
                          igNodeList (aliasTbl,movesTbl,moveAttrTbl)
                  val attr = valOf(Table.look(igNodeAttrTbl,u))
                      handle e => raise e
                  val degree = degreeBase cEnv igNodeAttrTbl
              in if attr <> PRECOLORED andalso 
                    moveRelated u = false andalso
                    degree(u) < K
                 then Table.enter(igNodeAttrTbl,u,SIMPLIFYWORK)
                 else igNodeAttrTbl
              end
          fun ok(t,r) =
              let val degree = degreeBase cEnv igNodeAttrTbl
                  val attr = valOf(Table.look(igNodeAttrTbl,t))
                      handle e => raise e
                  val tAdjNodes = Graph.adj(t)
              in degree(t) < K orelse
                 attr = PRECOLORED orelse
                 List.exists (fn n => Graph.eq(n,r)) tAdjNodes
              end
          fun conservative nodes =
              let val degree = degreeBase cEnv igNodeAttrTbl
                  val k = foldl (fn (n,k) => if degree(n) >= K
                                             then k + 1 else k) 
                                0 nodes
              in k < K
              end
          fun combine(aliasTbl,moveAttrTbl,igNodeAttrTbl,u,v) =
              (* first, update attribute for v *)
              let val igNodeAttrTbl' = 
                      let val attr = valOf(Table.look(igNodeAttrTbl,v))
                              handle e => raise e
                      in case attr
                          of FREEZEWORK => Table.enter(igNodeAttrTbl,
                                                       v,COALESCED)
                           | SPILLWORK => Table.enter(igNodeAttrTbl,
                                                      v,COALESCED)
                           | _ => raise IllegalState
                      end
                  val aliasTbl' = Table.enter(aliasTbl,v,u)
                  (* Added as per the errata for mciML *)
                  val moveAttrTbl' = enableMoves (movesTbl,moveAttrTbl) [v]

                  (* next, update attribute for adjacent node of v *)
                  val (igNodeAttrTbl',moveAttrTbl') = 
                      let val adjNodeList = NodeSet.listItems(getAdjacentNodesBase 
                                                                  cEnv igNodeAttrTbl' v)
                          val _ = app (fn n => Graph.mk_edge {from=n,to=u}) adjNodeList
                          val decrementDegree = decrementDegreeBase cEnv aliasTbl'
                      in foldl decrementDegree (igNodeAttrTbl',moveAttrTbl') adjNodeList
                      end

                  (* finally, update attribute for u *)
                  val degree' = degreeBase cEnv igNodeAttrTbl'
                  val igNodeAttrTbl' = 
                      let val attr = valOf(Table.look(igNodeAttrTbl',u))
                              handle e => raise e
                      in if attr = FREEZEWORK andalso degree'(u) >= K 
                         then Table.enter(igNodeAttrTbl',u,SPILLWORK)
                         else igNodeAttrTbl'
                      end

              in (aliasTbl',moveAttrTbl',igNodeAttrTbl')
              end
          fun coalesce0 (aliasTbl,moveAttrTbl,igNodeAttrTbl)
                        {fgNode=moveNode,duPair=(defNode,useNode)} = 
              let val x = getAlias(aliasTbl,defNode)
                  val y = getAlias(aliasTbl,useNode)
                  val moveAttr = valOf(Table.look(moveAttrTbl,moveNode))
                      handle e => raise e

                  val getAdjacentTemps = 
                      getAdjacentTempsBase cEnv igNodeAttrTbl
                  val getAdjacentNodes = getAdjacentNodesBase cEnv igNodeAttrTbl
                  fun okForAllAdj(u,v) = 
                      let val adjNodeList = 
                              NodeSet.listItems(getAdjacentNodes v)
                      in List.all (fn t => ok(t,u)) adjNodeList
                      end
                  fun conservativeForUVAdj(u,v) =
                      let val adjUVNodeSet = NodeSet.union(getAdjacentNodes u,
                                                           getAdjacentNodes v)
                          val adjUVNodeList = NodeSet.listItems adjUVNodeSet
                      in conservative(adjUVNodeList)
                      end
                  fun updateAttr(u,v) =
                      let val uAttr = valOf(Table.look(igNodeAttrTbl,u))
                          val vAttr = valOf(Table.look(igNodeAttrTbl,v))
                              handle e => raise e
                      in if Graph.eq(u,v) 
                         then let val moveAttrTbl' = 
                                      Table.enter(moveAttrTbl,moveNode,
                                                  COALESCED_MOVE)
                              in (aliasTbl,
                                  moveAttrTbl',
                                  addWorkList(moveAttrTbl',igNodeAttrTbl,u),
                                  selectStack)
                              end
                         else if vAttr = PRECOLORED orelse 
                                 List.exists (fn n => Graph.eq(n,u)) 
                                             (Graph.adj(v))
                         then let val moveAttrTbl' =
                                      Table.enter(moveAttrTbl,moveNode,
                                                  CONSTRAINED_MOVE)
                                  val igNodeAttrTbl' =
                                      addWorkList(moveAttrTbl',igNodeAttrTbl,u)
                              in (aliasTbl,
                                  moveAttrTbl',
                                  addWorkList(moveAttrTbl',igNodeAttrTbl',v),
                                  selectStack)
                              end
                         else if (uAttr = PRECOLORED andalso 
                                  okForAllAdj(u,v)) 
                                 orelse
                                 (uAttr <> PRECOLORED andalso
                                  conservativeForUVAdj(u,v))
                         then 
                             let val (aliasTbl',moveAttrTbl',igNodeAttrTbl') =
                                     combine (aliasTbl,
                                              Table.enter(moveAttrTbl,
                                                          moveNode,
                                                          COALESCED_MOVE),
                                              igNodeAttrTbl,u,v)
                             in (aliasTbl',
                                 moveAttrTbl',
                                 addWorkList(moveAttrTbl',igNodeAttrTbl',u),
                                 selectStack)
                             end
                         else (aliasTbl,
                               Table.enter(moveAttrTbl,moveNode,ACTIVE_MOVE),
                               igNodeAttrTbl,
                               selectStack)
                      end
              in if valOf(Table.look(igNodeAttrTbl,y)) = PRECOLORED
                 then updateAttr(y,x)
                 else updateAttr(x,y)
              end
      in coalesce0 (aliasTbl,moveAttrTbl,igNodeAttrTbl) (hd(workListMoves))
      end

  fun freezeMoves (cEnv as ENV{igNodeList,movesTbl,moves,registers,...})
                  (aliasTbl,moveAttrTbl,igNodeAttrTbl) u =
      let fun getMoveList moveNode = 
              List.filter (fn {fgNode,...} => 
                              Graph.eq(fgNode,moveNode)) moves 
          val nodeMovesU = 
              List.concat(List.map getMoveList
                                   (getOpt(Table.look(movesTbl,u),nil)))
          fun freezeMoves0 ({fgNode=m,duPair=(x,y)},
                            (aliasTbl,moveAttrTbl,igNodeAttrTbl)) = 
              let val K = length(registers)
                  val degree = degreeBase cEnv igNodeAttrTbl
                  val moveAttrTbl' = 
                      Table.enter(moveAttrTbl,m,FROZEN_MOVE)
                  val v = if Graph.eq(getAlias(aliasTbl,y),
                                      getAlias(aliasTbl,u))
                          then getAlias(aliasTbl,x)
                          else getAlias(aliasTbl,y)
                  val moveRelated = 
                      moveRelatedBase igNodeList (aliasTbl,
                                                  movesTbl,moveAttrTbl')
              in if moveRelated(v) = false andalso degree(v) < K 
                 then (aliasTbl,moveAttrTbl',
                       Table.enter(igNodeAttrTbl,v,SIMPLIFYWORK))
                 else (aliasTbl,moveAttrTbl',igNodeAttrTbl)
              end
      in foldl freezeMoves0 (aliasTbl,moveAttrTbl,igNodeAttrTbl) nodeMovesU
      end

  fun freeze cEnv (aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack) 
             freezeWorkList =
      let val u = hd(freezeWorkList)
          val (aliasTbl,moveAttrTbl,igNodeAttrTbl) =
              freezeMoves cEnv (aliasTbl,moveAttrTbl,
                                Table.enter(igNodeAttrTbl,u,SIMPLIFYWORK)) u
      in (aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack)
      end

  fun selectSpill cEnv (aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack) 
                  spillWorkList =
      let (* may introduce appropriate heuristic later. see mciML p.248 *)
          val m = hd(spillWorkList)
          val (aliasTbl,moveAttrTbl,igNodeAttrTbl) =
              freezeMoves cEnv (aliasTbl,moveAttrTbl,
                                Table.enter(igNodeAttrTbl,m,SIMPLIFYWORK)) m
      in (aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack)
      end

  fun color {interference,initial,spillCost,registers,
             moveAttrTbl=moveAttrTblInitial, (* Mutable *)
             movesTbl} =                     (* XXX Immutable??? *)
      let val allocTable = Temp.Table.empty (* Mutable *)
          val spillList = nil               (* Mutable *)
          val Liveness.IGRAPH{graph=iGraph, (* Immutable - Used once *)
                              gtemp,        (* Immutable - Moderately used *)
                              moves,        (* Immutable - Used in nodeMoves *)
                              tnode} =      (* Immutable - Moderately used? *)
              interference 
          val igNodeList = Graph.nodes(iGraph) (* Immutable - Heaviliy used *)

          val cEnv = ENV{movesTbl=movesTbl,igNodeList=igNodeList,moves=moves,
                         gtemp=gtemp,tnode=tnode,registers=registers}

          val igNodeAttrTblInitial = getInitialIgNodeAttrTbl(gtemp, initial, igNodeList)
          val aliasTbl = Table.empty (* Mutable *)
          val igNodeAttrTbl = 
              makeWorkList cEnv (igNodeAttrTblInitial,moveAttrTblInitial)
      in (allocTable, spillList)
      end
end