    -- | Graph Coloring.
    --      This is a generic graph coloring library, abstracted over the type of
    --      the node keys, nodes and colors.
    --

    module GraphColor (
            module GraphBase,
            module GraphOps,
            module GraphPpr,
            colorGraph
    )

    where

    import GraphBase
    import GraphOps
    import GraphPpr

    import Unique
    import UniqFM
    import UniqSet
    import Outputable

    import Data.Maybe
    import Data.List

    -- | Try to color a graph with this set of colors.
    --      Uses Chaitin's algorithm to color the graph.
    --      The graph is scanned for nodes which are deamed 'trivially colorable'. These nodes
    --      are pushed onto a stack and removed from the graph.
    --      Once this process is complete the graph can be colored by removing nodes from
    --      the stack (ie in reverse order) and assigning them colors different to their neighbors.
    --
    colorGraph
            :: ( Uniquable  k, Uniquable cls,  Uniquable  color
               , Eq color, Eq cls, Ord k
               , Outputable k, Outputable cls, Outputable color)
            => Bool                         -- ^ whether to do iterative coalescing
            -> Int                          -- ^ how many times we've tried to color this graph so far.
            -> UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
            -> Triv   k cls color           -- ^ fn to decide whether a node is trivially colorable.
            -> (Graph k cls color -> k)     -- ^ fn to choose a node to potentially leave uncolored if nothing is trivially colorable.
            -> Graph  k cls color           -- ^ the graph to color.

            -> ( Graph k cls color          -- the colored graph.
               , UniqSet k                  -- the set of nodes that we couldn't find a color for.
               , UniqFM  k )                -- map of regs (r1 -> r2) that were coaleced
                                            --       r1 should be replaced by r2 in the source

    colorGraph iterative spinCount colors triv spill graph0
     = let
            -- If we're not doing iterative coalescing then do an aggressive coalescing first time
            --      around and then conservative coalescing for subsequent passes.
            --
            --      Aggressive coalescing is a quick way to get rid of many reg-reg moves. However, if
            --      there is a lot of register pressure and we do it on every round then it can make the
            --      graph less colorable and prevent the algorithm from converging in a sensible number
            --      of cycles.
            --
            (graph_coalesced, kksCoalesce1)
             = if iterative
                    then (graph0, [])
                    else if spinCount == 0
                            then coalesceGraph True  triv graph0
                            else coalesceGraph False triv graph0

            -- run the scanner to slurp out all the trivially colorable nodes
            --      (and do coalescing if iterative coalescing is enabled)
            (ksTriv, ksProblems, kksCoalesce2)
                    = colorScan iterative triv spill graph_coalesced

            -- If iterative coalescing is enabled, the scanner will coalesce the graph as does its business.
            --      We need to apply all the coalescences found by the scanner to the original
            --      graph before doing assignColors.
            --
            --      Because we've got the whole, non-pruned graph here we turn on aggressive coalecing
            --      to force all the (conservative) coalescences found during scanning.
            --
            (graph_scan_coalesced, _)
                    = mapAccumL (coalesceNodes True triv) graph_coalesced kksCoalesce2

            -- color the trivially colorable nodes
            --      during scanning, keys of triv nodes were added to the front of the list as they were found
            --      this colors them in the reverse order, as required by the algorithm.
            (graph_triv, ksNoTriv)
                    = assignColors colors graph_scan_coalesced ksTriv

            -- try and color the problem nodes
            --      problem nodes are the ones that were left uncolored because they weren't triv.
            --      theres a change we can color them here anyway.
            (graph_prob, ksNoColor)
                    = assignColors colors graph_triv ksProblems

            -- if the trivially colorable nodes didn't color then something is probably wrong
            --      with the provided triv function.
            --
       in   if not $ null ksNoTriv
             then   pprPanic "colorGraph: trivially colorable nodes didn't color!" -- empty
                            (  empty
                            $$ text "ksTriv    = " <> ppr ksTriv
                            $$ text "ksNoTriv  = " <> ppr ksNoTriv
                            $$ text "colors    = " <> ppr colors
                            $$ empty
                            $$ dotGraph (\_ -> text "white") triv graph_triv)

             else   ( graph_prob
                    , mkUniqSet ksNoColor   -- the nodes that didn't color (spills)
                    , if iterative
                            then (listToUFM kksCoalesce2)
                            else (listToUFM kksCoalesce1))

    -- | Scan through the conflict graph separating out trivially colorable and
    --      potentially uncolorable (problem) nodes.
    --
    --      Checking whether a node is trivially colorable or not is a resonably expensive operation,
    --      so after a triv node is found and removed from the graph it's no good to return to the 'start'
    --      of the graph and recheck a bunch of nodes that will probably still be non-trivially colorable.
    --
    --      To ward against this, during each pass through the graph we collect up a list of triv nodes
    --      that were found, and only remove them once we've finished the pass. The more nodes we can delete
    --      at once the more likely it is that nodes we've already checked will become trivially colorable
    --      for the next pass.
    --
    --      TODO:   add work lists to finding triv nodes is easier.
    --              If we've just scanned the graph, and removed triv nodes, then the only
    --              nodes that we need to rescan are the ones we've removed edges from.

    colorScan
            :: ( Uniquable k, Uniquable cls, Uniquable color
               , Ord k,       Eq cls
               , Outputable k, Outputable cls)
            => Bool                         -- ^ whether to do iterative coalescing
            -> Triv k cls color             -- ^ fn to decide whether a node is trivially colorable
            -> (Graph k cls color -> k)     -- ^ fn to choose a node to potentially leave uncolored if nothing is trivially colorable.
            -> Graph k cls color            -- ^ the graph to scan

            -> ([k], [k], [(k, k)])         --  triv colorable nodes, problem nodes, pairs of nodes to coalesce

    colorScan iterative triv spill graph
            = colorScan_spin iterative triv spill graph [] [] []

    colorScan_spin
            :: ( Uniquable k, Uniquable cls, Uniquable color
               , Ord k,       Eq cls
               , Outputable k, Outputable cls)
            => Bool
            -> Triv k cls color
            -> (Graph k cls color -> k)
            -> Graph k cls color
            -> [k]
            -> [k]
            -> [(k, k)]
            -> ([k], [k], [(k, k)])

    colorScan_spin iterative triv spill graph
            ksTriv ksSpill kksCoalesce

            -- if the graph is empty then we're done
            | isNullUFM $ graphMap graph
            = (ksTriv, ksSpill, reverse kksCoalesce)

            -- Simplify:
            --      Look for trivially colorable nodes.
            --      If we can find some then remove them from the graph and go back for more.
            --
            | nsTrivFound@(_:_)
                    <-  scanGraph   (\node -> triv  (nodeClass node) (nodeConflicts node) (nodeExclusions node)

                                      -- for iterative coalescing we only want non-move related
                                      --    nodes here
                                      && (not iterative || isEmptyUniqSet (nodeCoalesce node)))
                            $ graph

            , ksTrivFound   <- map nodeId nsTrivFound
            , graph2        <- foldr (\k g -> let Just g' = delNode k g
                                              in  g')
                                    graph ksTrivFound

            = colorScan_spin iterative triv spill graph2
                    (ksTrivFound ++ ksTriv)
                    ksSpill
                    kksCoalesce

            -- Coalesce:
            --      If we're doing iterative coalescing and no triv nodes are avaliable
            --      then it's time for a coalescing pass.
            | iterative
            = case coalesceGraph False triv graph of

                    -- we were able to coalesce something
                    --      go back to Simplify and see if this frees up more nodes to be trivially colorable.
                    (graph2, kksCoalesceFound @(_:_))
                     -> colorScan_spin iterative triv spill graph2
                            ksTriv ksSpill (reverse kksCoalesceFound ++ kksCoalesce)

                    -- Freeze:
                    -- nothing could be coalesced (or was triv),
                    --      time to choose a node to freeze and give up on ever coalescing it.
                    (graph2, [])
                     -> case freezeOneInGraph graph2 of

                            -- we were able to freeze something
                            --      hopefully this will free up something for Simplify
                            (graph3, True)
                             -> colorScan_spin iterative triv spill graph3
                                    ksTriv ksSpill kksCoalesce

                            -- we couldn't find something to freeze either
                            --      time for a spill
                            (graph3, False)
                             -> colorScan_spill iterative triv spill graph3
                                    ksTriv ksSpill kksCoalesce

            -- spill time
            | otherwise
            = colorScan_spill iterative triv spill graph
                    ksTriv ksSpill kksCoalesce

    -- Select:
    -- we couldn't find any triv nodes or things to freeze or coalesce,
    --      and the graph isn't empty yet.. We'll have to choose a spill
    --      candidate and leave it uncolored.
    --
    colorScan_spill
            :: ( Uniquable k, Uniquable cls, Uniquable color
               , Ord k,       Eq cls
               , Outputable k, Outputable cls)
            => Bool
            -> Triv k cls color
            -> (Graph k cls color -> k)
            -> Graph k cls color
            -> [k]
            -> [k]
            -> [(k, k)]
            -> ([k], [k], [(k, k)])

    colorScan_spill iterative triv spill graph
            ksTriv ksSpill kksCoalesce

     = let  kSpill          = spill graph
            Just graph'     = delNode kSpill graph
       in   colorScan_spin iterative triv spill graph'
                    ksTriv (kSpill : ksSpill) kksCoalesce

    -- | Try to assign a color to all these nodes.

    assignColors
            :: ( Uniquable k, Uniquable cls, Uniquable color
               , Eq color, Outputable cls)
            => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
            -> Graph k cls color            -- ^ the graph
            -> [k]                          -- ^ nodes to assign a color to.
            -> ( Graph k cls color          -- the colored graph
               , [k])                       -- the nodes that didn't color.

    assignColors colors graph ks
            = assignColors' colors graph [] ks

     where  assignColors' _ graph prob []
                    = (graph, prob)

            assignColors' colors graph prob (k:ks)
             = case assignColor colors k graph of

                    -- couldn't color this node
                    Nothing         -> assignColors' colors graph (k : prob) ks

                    -- this node colored ok, so do the rest
                    Just graph'     -> assignColors' colors graph' prob ks

            assignColor colors u graph
                    | Just c        <- selectColor colors graph u
                    = Just (setColor u c graph)

                    | otherwise
                    = Nothing

    -- | Select a color for a certain node
    --      taking into account preferences, neighbors and exclusions.
    --      returns Nothing if no color can be assigned to this node.
    --
    selectColor
            :: ( Uniquable k, Uniquable cls, Uniquable color
               , Eq color, Outputable cls)
            => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
            -> Graph k cls color            -- ^ the graph
            -> k                            -- ^ key of the node to select a color for.
            -> Maybe color

    selectColor colors graph u
     = let  -- lookup the node
            Just node       = lookupNode graph u

            -- lookup the available colors for the class of this node.
            colors_avail
             = case lookupUFM colors (nodeClass node) of
                    Nothing -> pprPanic "selectColor: no colors available for class " (ppr (nodeClass node))
                    Just cs -> cs

            -- find colors we can't use because they're already being used
            --      by a node that conflicts with this one.
            Just nsConflicts
                            = sequence
                            $ map (lookupNode graph)
                            $ uniqSetToList
                            $ nodeConflicts node

            colors_conflict = mkUniqSet
                            $ catMaybes
                            $ map nodeColor nsConflicts

            -- the prefs of our neighbors
            colors_neighbor_prefs
                            = mkUniqSet
                            $ concat $ map nodePreference nsConflicts

            -- colors that are still valid for us
            colors_ok_ex    = minusUniqSet colors_avail (nodeExclusions node)
            colors_ok       = minusUniqSet colors_ok_ex colors_conflict

            -- the colors that we prefer, and are still ok
            colors_ok_pref  = intersectUniqSets
                                    (mkUniqSet $ nodePreference node) colors_ok

            -- the colors that we could choose while being nice to our neighbors
            colors_ok_nice  = minusUniqSet
                                    colors_ok colors_neighbor_prefs

            -- the best of all possible worlds..
            colors_ok_pref_nice
                            = intersectUniqSets
                                    colors_ok_nice colors_ok_pref

            -- make the decision
            chooseColor

                    -- everyone is happy, yay!
                    | not $ isEmptyUniqSet colors_ok_pref_nice
                    , c : _         <- filter (\x -> elementOfUniqSet x colors_ok_pref_nice)
                                            (nodePreference node)
                    = Just c

                    -- we've got one of our preferences
                    | not $ isEmptyUniqSet colors_ok_pref
                    , c : _         <- filter (\x -> elementOfUniqSet x colors_ok_pref)
                                            (nodePreference node)
                    = Just c

                    -- it wasn't a preference, but it was still ok
                    | not $ isEmptyUniqSet colors_ok
                    , c : _         <- uniqSetToList colors_ok
                    = Just c

                    -- no colors were available for us this time.
                    --      looks like we're going around the loop again..
                    | otherwise
                    = Nothing

       in   chooseColor