compiler/nativeGen/GraphColor.hs

   1
   2 -- | Graph Coloring.
   3 --      This is a generic graph coloring library, abstracted over the type of
   4 --      the node keys, nodes and colors.
   5 --
   6
   7 {-# OPTIONS_GHC -w #-}
   8 -- The above warning supression flag is a temporary kludge.
   9 -- While working on this module you are encouraged to remove it and fix
  10 -- any warnings in the module. See
  11 --     http://hackage.haskell.org/trac/ghc/wiki/WorkingConventions#Warnings
  12 -- for details
  13
  14 module GraphColor (
  15         module GraphBase,
  16         module GraphOps,
  17         module GraphPpr,
  18         colorGraph
  19 )
  20
  21 where
  22
  23 import GraphBase
  24 import GraphOps
  25 import GraphPpr
  26
  27 import Unique
  28 import UniqFM
  29 import UniqSet
  30 import Outputable
  31
  32 import Data.Maybe
  33 import Data.List
  34
  35
  36 -- | Try to color a graph with this set of colors.
  37 --      Uses Chaitin's algorithm to color the graph.
  38 --      The graph is scanned for nodes which are deamed 'trivially colorable'. These nodes
  39 --      are pushed onto a stack and removed from the graph.
  40 --      Once this process is complete the graph can be colored by removing nodes from
  41 --      the stack (ie in reverse order) and assigning them colors different to their neighbors.
  42 --
  43 colorGraph
  44         :: ( Uniquable  k, Uniquable cls,  Uniquable  color, Eq color
  45            , Outputable k, Outputable cls, Outputable color)
  46         => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
  47         -> Triv   k cls color           -- ^ fn to decide whether a node is trivially colorable.
  48         -> (Graph k cls color -> k)     -- ^ fn to choose a node to potentially leave uncolored if nothing is trivially colorable.
  49         -> Graph  k cls color           -- ^ the graph to color.
  50         -> ( Graph k cls color          -- ^ the colored graph.
  51            , UniqSet k )                -- ^ the set of nodes that we couldn't find a color for.
  52
  53 colorGraph colors triv spill graph0
  54  = let  -- run the scanner to slurp out all the trivially colorable nodes
  55         (ksTriv, ksProblems)    = colorScan colors triv spill [] emptyUniqSet graph0
  56
  57         -- color the trivially colorable nodes
  58         (graph1, ksNoTriv)      = assignColors colors graph0 ksTriv
  59
  60         -- try and color the problem nodes
  61         (graph2, ksNoColor)     = assignColors colors graph1 (uniqSetToList ksProblems)
  62
  63         -- if the trivially colorable nodes didn't color then something is wrong
  64         --      with the provided triv function.
  65    in   if not $ null ksNoTriv
  66          then   pprPanic "colorGraph: trivially colorable nodes didn't color!" empty
  67 {-                      (  empty
  68                         $$ text "ksTriv    = " <> ppr ksTriv
  69                         $$ text "ksNoTriv  = " <> ppr ksNoTriv
  70                         $$ empty
  71                         $$ dotGraph (\x -> text "white") triv graph1) -}
  72          else   (graph2, mkUniqSet ksNoColor)
  73
  74
  75 colorScan colors triv spill safe prob graph
  76
  77         -- empty graphs are easy to color.
  78         | isNullUFM $ graphMap graph
  79         = (safe, prob)
  80
  81         -- Try and find a trivially colorable node.
  82         | Just node     <- find (\node -> triv  (nodeClass node)
  83                                                 (nodeConflicts node)
  84                                                 (nodeExclusions node))
  85                                 $ eltsUFM $ graphMap graph
  86         , k             <- nodeId node
  87         = colorScan colors triv spill
  88                 (k : safe) prob (delNode k graph)
  89
  90         -- There was no trivially colorable node,
  91         --      Choose one to potentially leave uncolored. We /might/ be able to find
  92         --      a color for this later on, but no guarantees.
  93         | k             <- spill graph
  94         = colorScan colors triv spill
  95                 safe (addOneToUniqSet prob k) (delNode k graph)
  96
  97
  98
  99 -- | Try to assign a color to all these nodes.
 100
 101 assignColors
 102         :: ( Uniquable k, Uniquable cls, Uniquable color, Eq color )
 103         => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
 104         -> Graph k cls color            -- ^ the graph
 105         -> [k]                          -- ^ nodes to assign a color to.
 106         -> ( Graph k cls color          -- the colored graph
 107            , [k])                       -- the nodes that didn't color.
 108
 109 assignColors colors graph ks
 110         = assignColors' colors graph [] ks
 111
 112  where  assignColors' colors graph prob []
 113                 = (graph, prob)
 114
 115         assignColors' colors graph prob (k:ks)
 116          = case assignColor colors k graph of
 117
 118                 -- couldn't color this node
 119                 Nothing         -> assignColors' colors graph (k : prob) ks
 120
 121                 -- this node colored ok, so do the rest
 122                 Just graph'     -> assignColors' colors graph' prob ks
 123
 124
 125         assignColor colors u graph
 126                 | Just c        <- selectColor colors graph u
 127                 = Just (setColor u c graph)
 128
 129                 | otherwise
 130                 = Nothing
 131
 132
 133
 134 -- | Select a color for a certain node
 135 --      taking into account preferences, neighbors and exclusions.
 136 --      returns Nothing if no color can be assigned to this node.
 137 --
 138 --      TODO: avoid using the prefs of the neighbors, if at all possible.
 139 --
 140 selectColor
 141         :: ( Uniquable k, Uniquable cls, Uniquable color, Eq color)
 142         => UniqFM (UniqSet color)       -- ^ map of (node class -> set of colors available for this class).
 143         -> Graph k cls color            -- ^ the graph
 144         -> k                            -- ^ key of the node to select a color for.
 145         -> Maybe color
 146
 147 selectColor colors graph u
 148  = let  -- lookup the node
 149         Just node       = lookupNode graph u
 150
 151         -- lookup the available colors for the class of this node.
 152         Just colors_avail
 153                         = lookupUFM colors (nodeClass node)
 154
 155         -- colors we can't use because they're already being used
 156         --      by a node that conflicts with this one.
 157         Just nsConflicts
 158                         = sequence
 159                         $ map (lookupNode graph)
 160                         $ uniqSetToList
 161                         $ nodeConflicts node
 162
 163         colors_conflict = mkUniqSet
 164                         $ catMaybes
 165                         $ map nodeColor nsConflicts
 166
 167         -- colors that are still ok
 168         colors_ok_ex    = minusUniqSet colors_avail (nodeExclusions node)
 169         colors_ok       = minusUniqSet colors_ok_ex colors_conflict
 170
 171         -- the colors that we prefer, and are still ok
 172         colors_ok_pref  = intersectUniqSets
 173                                 (mkUniqSet $ nodePreference node) colors_ok
 174
 175         -- make the decision
 176         chooseColor
 177
 178                 -- we got one of our preferences, score!
 179                 | not $ isEmptyUniqSet colors_ok_pref
 180                 , c : rest      <- uniqSetToList colors_ok_pref
 181                 = Just c
 182
 183                 -- it wasn't a preference, but it was still ok
 184                 | not $ isEmptyUniqSet colors_ok
 185                 , c : rest      <- uniqSetToList colors_ok
 186                 = Just c
 187
 188                 -- leave this node uncolored
 189                 | otherwise
 190                 = Nothing
 191
 192    in   chooseColor
 193
 194
 195