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 T = Graph.Table
  structure S = IntListSet
  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 NS = ListSetFn(NodeKey) (* NodeSet *)

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

  (* For "coalesed" nodes, returns the node they are coalesced to
     (it assumes if a node has a aliasTbl entry, it's "coalesced").
     Otherwise, returns itself *)
  fun getAlias(aliasTbl,igNode) =
      case T.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(T.look(igNodeAttrTbl,igNode)) = attr
              handle e => raise e
      in List.filter hasAttr igNodeList
      end 

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

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

  fun nodeMoves igNodeList (aliasTbl,movesTbl,moveAttrTbl) n =
      let fun nodeMoves0 n =
              let val moveList = getOpt(T.look(movesTbl,n),nil)
                  fun isActiveMovesOrWorklistMoves moveNode =
                      let val moveAttr = valOf(T.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 T.empty for substitute here *)
          val moveRelated = 
              moveRelatedBase igNodeList (T.empty,movesTbl,moveAttrTbl)
          fun setAttr(igNode,igNodeAttrTbl) = 
              let val attr = valOf(T.look(igNodeAttrTbl,igNode))
                      handle e => raise e
                  val K = length(registers)
              in case attr
                  of INITIAL =>
                     if degree(igNode) >= K
                     then T.enter(igNodeAttrTbl,igNode,SPILLWORK)
                     else if moveRelated(igNode)
                     then T.enter(igNodeAttrTbl,igNode,FREEZEWORK)
                     else T.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(T.look(movesTbl,n),nil)
                  fun updateMoveAttrTbl(m,moveAttrTbl) =
                      let val attr = valOf(T.look(moveAttrTbl,m))
                              handle e => raise e
                      in if attr = ACTIVE_MOVE
                         then T.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 = 
                      NS.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 (T.enter(igNodeAttrTbl,m,FREEZEWORK),moveAttrTbl') 
                 else (T.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' = T.enter(igNodeAttrTbl,n,SELECTED)
          val selectStack' = n::selectStack
          val moveAttrTbl' = moveAttrTbl
          val adjNodeList = 
              NS.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(T.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 T.enter(igNodeAttrTbl,u,SIMPLIFYWORK)
                 else igNodeAttrTbl
              end
          fun ok(t,r) =
              let val degree = degreeBase cEnv igNodeAttrTbl
                  val attr = valOf(T.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(T.look(igNodeAttrTbl,v))
                              handle e => raise e
                      in case attr
                          of FREEZEWORK => T.enter(igNodeAttrTbl,v,COALESCED)
                           | SPILLWORK => T.enter(igNodeAttrTbl,v,COALESCED)
                           | _ => raise IllegalState
                      end
                  val aliasTbl' = T.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 = 
                              NS.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(T.look(igNodeAttrTbl',u))
                              handle e => raise e
                      in if attr = FREEZEWORK andalso degree'(u) >= K 
                         then T.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(T.look(moveAttrTbl,moveNode))
                      handle e => raise e
                  val getAdjacentNodes = 
                      getAdjacentNodesBase cEnv igNodeAttrTbl
                  fun okForAllAdj(u,v) = 
                      let val adjNodeList = NS.listItems(getAdjacentNodes v)
                      in List.all (fn t => ok(t,u)) adjNodeList
                      end
                  fun conservativeForUVAdj(u,v) =
                      let val adjUVNodeSet = NS.union(getAdjacentNodes u,
                                                      getAdjacentNodes v)
                          val adjUVNodeList = NS.listItems adjUVNodeSet
                      in conservative(adjUVNodeList)
                      end
                  fun updateAttr(u,v) =
                      let val uAttr = valOf(T.look(igNodeAttrTbl,u))
                          val vAttr = valOf(T.look(igNodeAttrTbl,v))
                              handle e => raise e
                      in if Graph.eq(u,v) 
                         then let val moveAttrTbl' = 
                                      T.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' =
                                      T.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,
                                              T.enter(moveAttrTbl,moveNode,
                                                      COALESCED_MOVE),
                                              igNodeAttrTbl,u,v)
                             in (aliasTbl',
                                 moveAttrTbl',
                                 addWorklist(moveAttrTbl',igNodeAttrTbl',u),
                                 selectStack)
                             end
                         else (aliasTbl,
                               T.enter(moveAttrTbl,moveNode,ACTIVE_MOVE),
                               igNodeAttrTbl,
                               selectStack)
                      end
              in if valOf(T.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(T.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' = T.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',
                       T.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,
                                T.enter(igNodeAttrTbl,u,SIMPLIFYWORK)) u
      in (aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack)
      end

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

  fun getWorklists (cEnv as ENV{igNodeList,moves,...}) 
                   (moveAttrTbl,igNodeAttrTbl) =
      let val getNodesWithAttr = getNodesWithAttrBase igNodeList igNodeAttrTbl 
          val simplifyWorklist = getNodesWithAttr SIMPLIFYWORK
          val worklistMoves = 
              let fun isWorklistMove {fgNode,duPair} =
                      valOf(T.look(moveAttrTbl,fgNode)) = WORKLIST_MOVE
              in List.filter isWorklistMove moves
              end
          val freezeWorklist = getNodesWithAttr FREEZEWORK
          val spillWorklist = getNodesWithAttr SPILLWORK
      in (simplifyWorklist,worklistMoves,freezeWorklist,spillWorklist)
      end
          
  fun allocLoop _ tables (nil,nil,nil,nil) = tables
    | allocLoop (cEnv as ENV{igNodeList,moves,...}) tables 
                (simplifyWorklist,worklistMoves,
                 freezeWorklist,spillWorklist) = 
      let val tables as (aliasTbl,moveAttrTbl,igNodeAttrTbl,selectStack) =
              if null(simplifyWorklist) = false
              then simplify cEnv tables simplifyWorklist
              else if null(worklistMoves) = false
              then coalesce cEnv tables worklistMoves
              else if null(freezeWorklist) = false
              then freeze cEnv tables freezeWorklist
              else if null(spillWorklist) = false
              then selectSpill cEnv tables spillWorklist
              else raise IllegalState
          val worklists = getWorklists cEnv (moveAttrTbl,igNodeAttrTbl)
      in allocLoop cEnv tables worklists
      end

  fun selectLoop _ _ igNodeAttrTbl colorTbl nil = (igNodeAttrTbl,colorTbl)
    | selectLoop (cEnv as ENV{registers,gtemp,...})
                   aliasTbl igNodeAttrTbl 
                   (colorTbl:Frame.register Graph.Table.table) (n::ns) = 
      let val w = Graph.adj(n)
          val okColors = 
              let fun takeAliasedColor(w,okColors) = 
                      let val wAlias = getAlias(aliasTbl,w)
                          val waAttr = valOf(T.look(igNodeAttrTbl,wAlias))
                          fun waColor() = 
                              valOf(T.look(colorTbl,wAlias))
                      in if (waAttr = COLORED orelse waAttr = PRECOLORED)
                            andalso S.member(okColors,waColor())
                         then S.delete(okColors,waColor())
                         else okColors
                      end
                  val okColorsInitial = S.addList(S.empty,registers)
              in foldl takeAliasedColor okColorsInitial w
              end
          fun setSpilledOrColored okColors igNodeAttrTbl colorTbl =
              if S.isEmpty okColors 
              then (T.enter(igNodeAttrTbl,n,SPILLED),colorTbl)
              else (T.enter(igNodeAttrTbl,n,COLORED),
                    T.enter(colorTbl,n,hd(S.listItems(okColors))))
          val (igNodeAttrTbl',colorTbl') = 
              setSpilledOrColored okColors igNodeAttrTbl colorTbl
      in selectLoop cEnv aliasTbl igNodeAttrTbl' colorTbl' ns
      end

  fun assignColors (cEnv as ENV{igNodeList,gtemp,...})
                   aliasTbl igNodeAttrTbl selectStack = 
      let val precoloredNodes = 
              getNodesWithAttrBase igNodeList igNodeAttrTbl PRECOLORED
          val colorTblInitial = 
              foldl (fn (n,table) => T.enter(table,n,gtemp n))
                    T.empty precoloredNodes
          val (igNodeAttrTbl,colorTbl) = 
              selectLoop cEnv aliasTbl igNodeAttrTbl 
                         colorTblInitial selectStack
          fun setColorForCoalesced(n,colorTbl) =
              let val attr = valOf(T.look(igNodeAttrTbl,n))
              in case attr
                      of COALESCED => 
                         let val aliasColor = 
                                 valOf(T.look(colorTbl,getAlias(aliasTbl,n)))
                         in T.enter(colorTbl,n,aliasColor)
                         end
                       | _ => colorTbl
              end
          val colorTbl = foldl setColorForCoalesced colorTbl igNodeList
      in (igNodeAttrTbl,colorTbl)
      end

  fun color {interference,initial,spillCost,registers,
             moveAttrTbl=moveAttrTblInitial,(* Mutable *)
             movesTbl} =                    (* XXX Immutable??? *)
      let 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 = T.empty (* Mutable *)
          val igNodeAttrTbl = 
              makeWorklist cEnv (igNodeAttrTblInitial,moveAttrTblInitial)
          val allocationTbl = Temp.Table.empty
      in (allocationTbl, spillList)
      end
end