-- | 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 GhcPrelude

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 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 coalesced
                                        --       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 available
        --      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
           , 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
           , 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)
                        $ nonDetEltsUniqSet
                        $ nodeConflicts node
                        -- See Note [Unique Determinism and code generation]

        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 : _         <- nonDetEltsUniqSet colors_ok
                -- See Note [Unique Determinism and code generation]
                = Just c

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

   in   chooseColor