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