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,
                         gtemp: Graph.node -> Temp.temp,
                         tnode: Temp.temp -> Graph.node}
  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,
                         gtemp: Graph.node -> Temp.temp,
                         tnode: Temp.temp -> Graph.node}

  val registers = ref (Frame.registers)
  val K = length(!registers)
  fun getInitialIgNodeAttrTbl(gtemp,initial,igNodeList) = 
      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,igNode) =
      Map.listItems(Map.filter (fn n => Graph.eq(n,igNode) = false andalso
                                        Graph.eq(getAlias(aliasTbl,n),igNode)) 
                               aliasTbl)

  fun getAdjacentTempsBase (cEnv as ENV{igNodeList,gtemp,...})
                           (aliasTbl,igNodeAttrTbl) igNode = 
      let fun getAdjacentTempsBase0 igNode = 
              let fun getTempsWithAttr(attr) = 
                      let fun hasAttr(igNode) =
                              (valOf(Graph.Table.look(igNodeAttrTbl,igNode)) 
                               = attr) 
                              handle e => raise e
                      in List.map gtemp (List.filter hasAttr igNodeList)
                      end 
                  fun makeNewSet(list) = Set.addList(Set.empty,list)
                  val adjIgNodes = Graph.adj(igNode)
                  val sTempSet = makeNewSet(getTempsWithAttr(SELECTED))
                  val scTempSet = Set.addList(sTempSet,
                                              getTempsWithAttr(COALESCED))
                  val adjTempSet = makeNewSet(List.map gtemp adjIgNodes)
              in Set.difference(adjTempSet,scTempSet)
              end
          val aliasList = getAliasNodeList(aliasTbl,igNode)
          val adjTempsSetList = map getAdjacentTempsBase0 (igNode::aliasList)
      in foldl Set.union Set.empty adjTempsSetList
      end

  fun degreeBase cEnv (aliasTbl,igNodeAttrTbl) igNode = 
      Set.numItems(getAdjacentTempsBase cEnv (aliasTbl,igNodeAttrTbl) igNode)

  fun nodeMoves (aliasTbl,movesTbl,moveAttrTbl) igNode =
      let fun nodeMoves0 igNode =
              let val moveList = getOpt(Table.look(movesTbl,igNode),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,igNode)
          val result = List.map nodeMoves0 (igNode::aliasList)
          (* caution - result may contain duplicate nodes *)
      in List.concat result
      end
  fun moveRelatedBase (aliasTbl,movesTbl,moveAttrTbl) igNode = 
      length(nodeMoves(aliasTbl,movesTbl,moveAttrTbl) igNode) <> 0

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

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

  fun enableMoves (movesTbl,moveAttrTbl) igNode =
      let val moveListN = getOpt(Table.look(movesTbl,igNode),nil)
          fun updateAttr(fgNode) = 
              let val attr = valOf(Table.look(moveAttrTbl,fgNode))
                      handle e => raise e
              in case attr 
                  of ACTIVE_MOVE => WORKLIST_MOVE
                   | _ => attr
              end
          fun updateMoveAttrTbl(fgNode,moveAttrTblOrig) =
              Table.enter(moveAttrTblOrig,fgNode,updateAttr(fgNode))
      in foldl updateMoveAttrTbl moveAttrTbl moveListN
      end

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

  fun simplifyOneNode (cEnv as ENV{tnode,...}) aliasTbl
                      (igNode,(igNodeAttrTbl,moveAttrTbl,selectStack)) =
      let val origAttr = valOf(Table.look(igNodeAttrTbl,igNode))
              handle e => raise e
          val igNodeAttrTbl' = Table.enter(igNodeAttrTbl,igNode,SELECTED)
          val selectStack' = igNode::selectStack
          val moveAttrTbl' = moveAttrTbl
          (* XXX should add node "m" to adjTempSet? *)
          val adjTempSet = 
              getAdjacentTempsBase cEnv (aliasTbl,igNodeAttrTbl') igNode
          val adjNodeList = List.map tnode (Set.listItems adjTempSet)
          val decrementDegree = decrementDegreeBase cEnv aliasTbl

          val (igNodeAttrTbl'',moveAttrTbl'') =
              foldl decrementDegree (igNodeAttrTbl',moveAttrTbl') adjNodeList
      in (igNodeAttrTbl'',moveAttrTbl'',selectStack')
      end

  fun isSimplifyWork igNodeAttrTbl igNode =
      (valOf(Graph.Table.look(igNodeAttrTbl,igNode)) = SIMPLIFYWORK) 
      handle e => raise e

  fun simplify cEnv aliasTbl attrs nil = attrs
    | simplify (cEnv as ENV{igNodeList,...})
               aliasTbl attrs simplifyWorkNodeList =
      let val attrs' as (igNodeAttrTbl',moveAttrTbl',selectStack') =
              foldl (simplifyOneNode cEnv aliasTbl) attrs simplifyWorkNodeList
          val simplifyWorkNodeList' = 
              List.filter (isSimplifyWork igNodeAttrTbl') igNodeList
      in simplify cEnv aliasTbl attrs' simplifyWorkNodeList'
      end

  fun coalesce (cEnv as ENV{igNodeList,movesTbl,tnode,...})
               (aliasTbl,moveAttrTbl,igNodeAttrTbl) moves =
      let val degree = degreeBase cEnv (aliasTbl,igNodeAttrTbl) 
          fun addWorkList(igNodeAttrTbl,igNode) = 
              let val moveRelated = moveRelatedBase 
                                        (aliasTbl,movesTbl,moveAttrTbl)
                  val attr = valOf(Table.look(igNodeAttrTbl,igNode))
                      handle e => raise e
              in if attr <> PRECOLORED andalso 
                    moveRelated igNode = false andalso
                    degree(igNode) < K - 1 
                 then Table.enter(igNodeAttrTbl,igNode,SIMPLIFYWORK)
                 else igNodeAttrTbl
              end
          fun ok(t,r) =
              let 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 igNodeList =
              let val k = foldl (fn (igNode,k) => if degree(igNode) >= K
                                                  then k + 1 else k) 
                                0 igNodeList
              in k < K
              end
          fun combine(aliasTbl,igNodeAttrTbl,u,v) =
              let fun updateAttrV igNode =
                      let val origAttr =
                              valOf(Table.look(igNodeAttrTbl,igNode))
                              handle e => raise e
                      in case origAttr
                          of FREEZEWORK => Table.enter(igNodeAttrTbl,
                                                       igNode,COALESCED)
                           | SPILLWORK => Table.enter(igNodeAttrTbl,
                                                      igNode,COALESCED)
                           | _ => raise IllegalState
                      end
                  val igNodeAttrTbl' = updateAttrV v
                  val aliasTbl' = Table.enter(aliasTbl,v,u)
                  val moveAttrTbl' = enableMoves (movesTbl,moveAttrTbl) v
                  val degree' = degreeBase cEnv (aliasTbl',igNodeAttrTbl) 
                  fun updateAttrU igNode =
                      let val origAttr =
                              valOf(Table.look(igNodeAttrTbl,igNode))
                              handle e => raise e
                      in if origAttr = FREEZEWORK andalso 
                            degree'(igNode) >= K 
                         then Table.enter(igNodeAttrTbl,igNode,SPILLWORK)
                         else igNodeAttrTbl
                      end
              in (aliasTbl',moveAttrTbl',igNodeAttrTbl')
              end
          fun coalesce0 (aliasTbl,moveAttrTbl,igNodeAttrTbl)
                        {moveNode,duPair=(defNode,useNode)} =
              let val x = getAlias(defNode)
                  val y = getAlias(useNode)
                  val moveAttr = 
                      valOf(Table.look(moveAttrTbl,moveNode))
                      handle e => raise e
                  val getAdjacentTemps = 
                      getAdjacentTempsBase cEnv (aliasTbl,igNodeAttrTbl)
                  fun okForAllAdj(u,v) = 
                      let val adjNodeList =
                              List.map tnode 
                                       (Set.listItems(getAdjacentTemps v))
                      in List.all (fn t => ok(t,u)) adjNodeList
                      end
                  fun conservativeForUVAdj(u,v) =
                      let val adjUVTempSet = Set.union(getAdjacentTemps u,
                                                       getAdjacentTemps v)
                          val adjUVNodeList = 
                              List.map tnode (Set.listItems adjUVTempSet)
                      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 (aliasTbl,
                               Table.enter(moveAttrTbl,moveNode,
                                           COALESCED_MOVE),
                               addWorkList(igNodeAttrTbl,u))
                         else if vAttr = PRECOLORED orelse 
                                 List.exists (fn n => Graph.eq(n,u)) 
                                             (Graph.adj(v))
                         then (aliasTbl,
                               Table.enter(moveAttrTbl,moveNode,
                                           CONSTRAINED_MOVE),
                               addWorkList(addWorkList(igNodeAttrTbl,u),v))
                         else if uAttr = PRECOLORED andalso 
                                 okForAllAdj(u,v) orelse
                                 uAttr <> PRECOLORED andalso
                                 conservativeForUVAdj(u,v)
                         then combine (aliasTbl,
                                       igNodeAttrTbl,u,v) (* XXX *)
                         else (aliasTbl,
                               Table.enter(moveAttrTbl,moveNode,ACTIVE_MOVE),
                               igNodeAttrTbl)
                      end
              in if valOf(Table.look(igNodeAttrTbl,y)) = PRECOLORED
                 then updateAttr(y,x)
                 else updateAttr(x,y)
              end
      in coalesce0 (aliasTbl,moveAttrTbl,igNodeAttrTbl) moves
      end

  fun color {interference,initial,spillCost,registers,
             moveAttrTbl=moveAttrTblInitial, (* Mutable *)
             movesTbl                        (* 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,
                         gtemp=gtemp,tnode=tnode}

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