move generic graph-colouring code into util
[ghc-hetmet.git] / compiler / utils / GraphColor.hs
diff --git a/compiler/utils/GraphColor.hs b/compiler/utils/GraphColor.hs
new file mode 100644 (file)
index 0000000..307803a
--- /dev/null
@@ -0,0 +1,332 @@
+
+-- | Graph Coloring.
+--     This is a generic graph coloring library, abstracted over the type of
+--     the node keys, nodes and colors.
+--
+{-# OPTIONS -fno-warn-missing-signatures #-}
+
+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
+       -> 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 colors triv spill graph0
+ = let
+       -- if we're not doing iterative coalescing, then just do a single coalescing
+       --      pass at the front. This won't be as good but should still eat up a
+       --      lot of the reg-reg moves.
+       (graph_coalesced, kksCoalesce1)
+               = if not iterative
+                       then coalesceGraph False triv graph0
+                       else (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
+                       $$ empty
+                       $$ dotGraph (\x -> text "white") triv graph1) -}
+
+        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 color)
+       => 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 iterative triv spill graph
+       ksTriv ksSpill kksCoalesce
+
+       -- if the graph is empty then we're done
+       | isNullUFM $ graphMap graph
+       = (ksTriv, ksSpill, 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
+       , graph3        <- foldr (\k g -> let Just g' = delNode k g
+                                         in  g')
+                               graph ksTrivFound
+
+       = colorScan_spin iterative triv spill graph3
+               (ksTrivFound ++ ksTriv)
+               ksSpill
+               kksCoalesce
+
+       -- Coalesce:
+       --      If we're doing iterative coalescing and no triv nodes are avaliable
+       --      then it's type for a coalescing pass.
+       | iterative
+       = case coalesceGraph False triv graph of
+
+               -- we were able to coalesce something
+               --      go back and see if this frees up more nodes to be trivially colorable.
+               (graph2, kksCoalesceFound @(_:_))
+                -> colorScan_spin iterative triv spill graph2
+                       ksTriv ksSpill (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 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 )
+       => 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)
+       => 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.
+       Just colors_avail
+                       = lookupUFM colors (nodeClass node)
+
+       -- 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 
+
+
+