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