-- | 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 :: Bool
-> Int
-> UniqFM (UniqSet color)
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> (Graph k cls color, UniqSet k, UniqFM k)
colorGraph Bool
iterative Int
spinCount UniqFM (UniqSet color)
colors Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
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 k cls color
graph_coalesced, [(k, k)]
kksCoalesce1)
         = if Bool
iterative
                then (Graph k cls color
graph0, [])
                else if Int
spinCount Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0
                        then Bool
-> Triv k cls color
-> Graph k cls color
-> (Graph k cls color, [(k, k)])
forall k cls color.
(Uniquable k, Ord k, Eq cls, Outputable k) =>
Bool
-> Triv k cls color
-> Graph k cls color
-> (Graph k cls color, [(k, k)])
coalesceGraph Bool
True  Triv k cls color
triv Graph k cls color
graph0
                        else Bool
-> Triv k cls color
-> Graph k cls color
-> (Graph k cls color, [(k, k)])
forall k cls color.
(Uniquable k, Ord k, Eq cls, Outputable k) =>
Bool
-> Triv k cls color
-> Graph k cls color
-> (Graph k cls color, [(k, k)])
coalesceGraph Bool
False Triv k cls color
triv Graph k cls color
graph0

        -- run the scanner to slurp out all the trivially colorable nodes
        --      (and do coalescing if iterative coalescing is enabled)
        ([k]
ksTriv, [k]
ksProblems, [(k, k)]
kksCoalesce2)
                = Bool
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> ([k], [k], [(k, k)])
forall k cls color.
(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)])
colorScan Bool
iterative Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
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 k cls color
graph_scan_coalesced, [Maybe (k, k)]
_)
                = (Graph k cls color -> (k, k) -> (Graph k cls color, Maybe (k, k)))
-> Graph k cls color
-> [(k, k)]
-> (Graph k cls color, [Maybe (k, k)])
forall (t :: * -> *) a b c.
Traversable t =>
(a -> b -> (a, c)) -> a -> t b -> (a, t c)
mapAccumL (Bool
-> Triv k cls color
-> Graph k cls color
-> (k, k)
-> (Graph k cls color, Maybe (k, k))
forall k cls color.
(Uniquable k, Ord k, Eq cls) =>
Bool
-> Triv k cls color
-> Graph k cls color
-> (k, k)
-> (Graph k cls color, Maybe (k, k))
coalesceNodes Bool
True Triv k cls color
triv) Graph k cls color
graph_coalesced [(k, k)]
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 k cls color
graph_triv, [k]
ksNoTriv)
                = UniqFM (UniqSet color)
-> Graph k cls color -> [k] -> (Graph k cls color, [k])
forall k cls color.
(Uniquable k, Uniquable cls, Uniquable color, Outputable cls) =>
UniqFM (UniqSet color)
-> Graph k cls color -> [k] -> (Graph k cls color, [k])
assignColors UniqFM (UniqSet color)
colors Graph k cls color
graph_scan_coalesced [k]
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 k cls color
graph_prob, [k]
ksNoColor)
                = UniqFM (UniqSet color)
-> Graph k cls color -> [k] -> (Graph k cls color, [k])
forall k cls color.
(Uniquable k, Uniquable cls, Uniquable color, Outputable cls) =>
UniqFM (UniqSet color)
-> Graph k cls color -> [k] -> (Graph k cls color, [k])
assignColors UniqFM (UniqSet color)
colors Graph k cls color
graph_triv [k]
ksProblems

        -- if the trivially colorable nodes didn't color then something is probably wrong
        --      with the provided triv function.
        --
   in   if Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ [k] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [k]
ksNoTriv
         then   String -> SDoc -> (Graph k cls color, UniqSet k, UniqFM k)
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"colorGraph: trivially colorable nodes didn't color!" -- empty
                        (  SDoc
empty
                        SDoc -> SDoc -> SDoc
$$ String -> SDoc
text String
"ksTriv    = " SDoc -> SDoc -> SDoc
<> [k] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [k]
ksTriv
                        SDoc -> SDoc -> SDoc
$$ String -> SDoc
text String
"ksNoTriv  = " SDoc -> SDoc -> SDoc
<> [k] -> SDoc
forall a. Outputable a => a -> SDoc
ppr [k]
ksNoTriv
                        SDoc -> SDoc -> SDoc
$$ String -> SDoc
text String
"colors    = " SDoc -> SDoc -> SDoc
<> UniqFM (UniqSet color) -> SDoc
forall a. Outputable a => a -> SDoc
ppr UniqFM (UniqSet color)
colors
                        SDoc -> SDoc -> SDoc
$$ SDoc
empty
                        SDoc -> SDoc -> SDoc
$$ (color -> SDoc) -> Triv k cls color -> Graph k cls color -> SDoc
forall k cls color.
(Uniquable k, Outputable k, Outputable cls, Outputable color) =>
(color -> SDoc) -> Triv k cls color -> Graph k cls color -> SDoc
dotGraph (\color
_ -> String -> SDoc
text String
"white") Triv k cls color
triv Graph k cls color
graph_triv)

         else   ( Graph k cls color
graph_prob
                , [k] -> UniqSet k
forall a. Uniquable a => [a] -> UniqSet a
mkUniqSet [k]
ksNoColor   -- the nodes that didn't color (spills)
                , if Bool
iterative
                        then ([(k, k)] -> UniqFM k
forall key elt. Uniquable key => [(key, elt)] -> UniqFM elt
listToUFM [(k, k)]
kksCoalesce2)
                        else ([(k, k)] -> UniqFM k
forall key elt. Uniquable key => [(key, elt)] -> UniqFM elt
listToUFM [(k, k)]
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 reasonably 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 :: Bool
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> ([k], [k], [(k, k)])
colorScan Bool
iterative Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
graph
        = Bool
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> [k]
-> [k]
-> [(k, k)]
-> ([k], [k], [(k, k)])
forall k cls color.
(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 Bool
iterative Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
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 :: Bool
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> [k]
-> [k]
-> [(k, k)]
-> ([k], [k], [(k, k)])
colorScan_spin Bool
iterative Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
graph
        [k]
ksTriv [k]
ksSpill [(k, k)]
kksCoalesce

        -- if the graph is empty then we're done
        | UniqFM (Node k cls color) -> Bool
forall elt. UniqFM elt -> Bool
isNullUFM (UniqFM (Node k cls color) -> Bool)
-> UniqFM (Node k cls color) -> Bool
forall a b. (a -> b) -> a -> b
$ Graph k cls color -> UniqFM (Node k cls color)
forall k cls color. Graph k cls color -> UniqFM (Node k cls color)
graphMap Graph k cls color
graph
        = ([k]
ksTriv, [k]
ksSpill, [(k, k)] -> [(k, k)]
forall a. [a] -> [a]
reverse [(k, k)]
kksCoalesce)

        -- Simplify:
        --      Look for trivially colorable nodes.
        --      If we can find some then remove them from the graph and go back for more.
        --
        | nsTrivFound :: [Node k cls color]
nsTrivFound@(Node k cls color
_:[Node k cls color]
_)
                <-  (Node k cls color -> Bool)
-> Graph k cls color -> [Node k cls color]
forall k cls color.
(Node k cls color -> Bool)
-> Graph k cls color -> [Node k cls color]
scanGraph   (\Node k cls color
node -> Triv k cls color
triv  (Node k cls color -> cls
forall k cls color. Node k cls color -> cls
nodeClass Node k cls color
node) (Node k cls color -> UniqSet k
forall k cls color. Node k cls color -> UniqSet k
nodeConflicts Node k cls color
node) (Node k cls color -> UniqSet color
forall k cls color. Node k cls color -> UniqSet color
nodeExclusions Node k cls color
node)

                                  -- for iterative coalescing we only want non-move related
                                  --    nodes here
                                  Bool -> Bool -> Bool
&& (Bool -> Bool
not Bool
iterative Bool -> Bool -> Bool
|| UniqSet k -> Bool
forall a. UniqSet a -> Bool
isEmptyUniqSet (Node k cls color -> UniqSet k
forall k cls color. Node k cls color -> UniqSet k
nodeCoalesce Node k cls color
node)))
                        (Graph k cls color -> [Node k cls color])
-> Graph k cls color -> [Node k cls color]
forall a b. (a -> b) -> a -> b
$ Graph k cls color
graph

        , [k]
ksTrivFound   <- (Node k cls color -> k) -> [Node k cls color] -> [k]
forall a b. (a -> b) -> [a] -> [b]
map Node k cls color -> k
forall k cls color. Node k cls color -> k
nodeId [Node k cls color]
nsTrivFound
        , Graph k cls color
graph2        <- (k -> Graph k cls color -> Graph k cls color)
-> Graph k cls color -> [k] -> Graph k cls color
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr (\k
k Graph k cls color
g -> let Just Graph k cls color
g' = k -> Graph k cls color -> Maybe (Graph k cls color)
forall k cls color.
Uniquable k =>
k -> Graph k cls color -> Maybe (Graph k cls color)
delNode k
k Graph k cls color
g
                                          in  Graph k cls color
g')
                                Graph k cls color
graph [k]
ksTrivFound

        = Bool
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> [k]
-> [k]
-> [(k, k)]
-> ([k], [k], [(k, k)])
forall k cls color.
(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 Bool
iterative Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
graph2
                ([k]
ksTrivFound [k] -> [k] -> [k]
forall a. [a] -> [a] -> [a]
++ [k]
ksTriv)
                [k]
ksSpill
                [(k, k)]
kksCoalesce

        -- Coalesce:
        --      If we're doing iterative coalescing and no triv nodes are available
        --      then it's time for a coalescing pass.
        | Bool
iterative
        = case Bool
-> Triv k cls color
-> Graph k cls color
-> (Graph k cls color, [(k, k)])
forall k cls color.
(Uniquable k, Ord k, Eq cls, Outputable k) =>
Bool
-> Triv k cls color
-> Graph k cls color
-> (Graph k cls color, [(k, k)])
coalesceGraph Bool
False Triv k cls color
triv Graph k cls color
graph of

                -- we were able to coalesce something
                --      go back to Simplify and see if this frees up more nodes to be trivially colorable.
                (Graph k cls color
graph2, kksCoalesceFound :: [(k, k)]
kksCoalesceFound@((k, k)
_:[(k, k)]
_))
                 -> Bool
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> [k]
-> [k]
-> [(k, k)]
-> ([k], [k], [(k, k)])
forall k cls color.
(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 Bool
iterative Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
graph2
                        [k]
ksTriv [k]
ksSpill ([(k, k)] -> [(k, k)]
forall a. [a] -> [a]
reverse [(k, k)]
kksCoalesceFound [(k, k)] -> [(k, k)] -> [(k, k)]
forall a. [a] -> [a] -> [a]
++ [(k, k)]
kksCoalesce)

                -- Freeze:
                -- nothing could be coalesced (or was triv),
                --      time to choose a node to freeze and give up on ever coalescing it.
                (Graph k cls color
graph2, [])
                 -> case Graph k cls color -> (Graph k cls color, Bool)
forall k cls color.
Uniquable k =>
Graph k cls color -> (Graph k cls color, Bool)
freezeOneInGraph Graph k cls color
graph2 of

                        -- we were able to freeze something
                        --      hopefully this will free up something for Simplify
                        (Graph k cls color
graph3, Bool
True)
                         -> Bool
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> [k]
-> [k]
-> [(k, k)]
-> ([k], [k], [(k, k)])
forall k cls color.
(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 Bool
iterative Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
graph3
                                [k]
ksTriv [k]
ksSpill [(k, k)]
kksCoalesce

                        -- we couldn't find something to freeze either
                        --      time for a spill
                        (Graph k cls color
graph3, Bool
False)
                         -> Bool
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> [k]
-> [k]
-> [(k, k)]
-> ([k], [k], [(k, k)])
forall k cls color.
(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 Bool
iterative Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
graph3
                                [k]
ksTriv [k]
ksSpill [(k, k)]
kksCoalesce

        -- spill time
        | Bool
otherwise
        = Bool
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> [k]
-> [k]
-> [(k, k)]
-> ([k], [k], [(k, k)])
forall k cls color.
(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 Bool
iterative Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
graph
                [k]
ksTriv [k]
ksSpill [(k, k)]
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 :: Bool
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> [k]
-> [k]
-> [(k, k)]
-> ([k], [k], [(k, k)])
colorScan_spill Bool
iterative Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
graph
        [k]
ksTriv [k]
ksSpill [(k, k)]
kksCoalesce

 = let  kSpill :: k
kSpill          = Graph k cls color -> k
spill Graph k cls color
graph
        Just Graph k cls color
graph'     = k -> Graph k cls color -> Maybe (Graph k cls color)
forall k cls color.
Uniquable k =>
k -> Graph k cls color -> Maybe (Graph k cls color)
delNode k
kSpill Graph k cls color
graph
   in   Bool
-> Triv k cls color
-> (Graph k cls color -> k)
-> Graph k cls color
-> [k]
-> [k]
-> [(k, k)]
-> ([k], [k], [(k, k)])
forall k cls color.
(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 Bool
iterative Triv k cls color
triv Graph k cls color -> k
spill Graph k cls color
graph'
                [k]
ksTriv (k
kSpill k -> [k] -> [k]
forall a. a -> [a] -> [a]
: [k]
ksSpill) [(k, k)]
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 :: UniqFM (UniqSet color)
-> Graph k cls color -> [k] -> (Graph k cls color, [k])
assignColors UniqFM (UniqSet color)
colors Graph k cls color
graph [k]
ks
        = UniqFM (UniqSet color)
-> Graph k cls color -> [k] -> [k] -> (Graph k cls color, [k])
forall a cls color.
(Uniquable a, Uniquable cls, Uniquable color, Outputable cls) =>
UniqFM (UniqSet color)
-> Graph a cls color -> [a] -> [a] -> (Graph a cls color, [a])
assignColors' UniqFM (UniqSet color)
colors Graph k cls color
graph [] [k]
ks

 where  assignColors' :: UniqFM (UniqSet color)
-> Graph a cls color -> [a] -> [a] -> (Graph a cls color, [a])
assignColors' UniqFM (UniqSet color)
_ Graph a cls color
graph [a]
prob []
                = (Graph a cls color
graph, [a]
prob)

        assignColors' UniqFM (UniqSet color)
colors Graph a cls color
graph [a]
prob (a
k:[a]
ks)
         = case UniqFM (UniqSet color)
-> a -> Graph a cls color -> Maybe (Graph a cls color)
forall k cls color.
(Uniquable k, Uniquable cls, Uniquable color, Outputable cls) =>
UniqFM (UniqSet color)
-> k -> Graph k cls color -> Maybe (Graph k cls color)
assignColor UniqFM (UniqSet color)
colors a
k Graph a cls color
graph of

                -- couldn't color this node
                Maybe (Graph a cls color)
Nothing         -> UniqFM (UniqSet color)
-> Graph a cls color -> [a] -> [a] -> (Graph a cls color, [a])
assignColors' UniqFM (UniqSet color)
colors Graph a cls color
graph (a
k a -> [a] -> [a]
forall a. a -> [a] -> [a]
: [a]
prob) [a]
ks

                -- this node colored ok, so do the rest
                Just Graph a cls color
graph'     -> UniqFM (UniqSet color)
-> Graph a cls color -> [a] -> [a] -> (Graph a cls color, [a])
assignColors' UniqFM (UniqSet color)
colors Graph a cls color
graph' [a]
prob [a]
ks


        assignColor :: UniqFM (UniqSet color)
-> k -> Graph k cls color -> Maybe (Graph k cls color)
assignColor UniqFM (UniqSet color)
colors k
u Graph k cls color
graph
                | Just color
c        <- UniqFM (UniqSet color) -> Graph k cls color -> k -> Maybe color
forall k cls color.
(Uniquable k, Uniquable cls, Uniquable color, Outputable cls) =>
UniqFM (UniqSet color) -> Graph k cls color -> k -> Maybe color
selectColor UniqFM (UniqSet color)
colors Graph k cls color
graph k
u
                = Graph k cls color -> Maybe (Graph k cls color)
forall a. a -> Maybe a
Just (k -> color -> Graph k cls color -> Graph k cls color
forall k color cls.
Uniquable k =>
k -> color -> Graph k cls color -> Graph k cls color
setColor k
u color
c Graph k cls color
graph)

                | Bool
otherwise
                = Maybe (Graph k cls color)
forall a. Maybe a
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 :: UniqFM (UniqSet color) -> Graph k cls color -> k -> Maybe color
selectColor UniqFM (UniqSet color)
colors Graph k cls color
graph k
u
 = let  -- lookup the node
        Just Node k cls color
node       = Graph k cls color -> k -> Maybe (Node k cls color)
forall k cls color.
Uniquable k =>
Graph k cls color -> k -> Maybe (Node k cls color)
lookupNode Graph k cls color
graph k
u

        -- lookup the available colors for the class of this node.
        colors_avail :: UniqSet color
colors_avail
         = case UniqFM (UniqSet color) -> cls -> Maybe (UniqSet color)
forall key elt. Uniquable key => UniqFM elt -> key -> Maybe elt
lookupUFM UniqFM (UniqSet color)
colors (Node k cls color -> cls
forall k cls color. Node k cls color -> cls
nodeClass Node k cls color
node) of
                Maybe (UniqSet color)
Nothing -> String -> SDoc -> UniqSet color
forall a. HasCallStack => String -> SDoc -> a
pprPanic String
"selectColor: no colors available for class " (cls -> SDoc
forall a. Outputable a => a -> SDoc
ppr (Node k cls color -> cls
forall k cls color. Node k cls color -> cls
nodeClass Node k cls color
node))
                Just UniqSet color
cs -> UniqSet color
cs

        -- find colors we can't use because they're already being used
        --      by a node that conflicts with this one.
        Just [Node k cls color]
nsConflicts
                        = [Maybe (Node k cls color)] -> Maybe [Node k cls color]
forall (t :: * -> *) (m :: * -> *) a.
(Traversable t, Monad m) =>
t (m a) -> m (t a)
sequence
                        ([Maybe (Node k cls color)] -> Maybe [Node k cls color])
-> [Maybe (Node k cls color)] -> Maybe [Node k cls color]
forall a b. (a -> b) -> a -> b
$ (k -> Maybe (Node k cls color))
-> [k] -> [Maybe (Node k cls color)]
forall a b. (a -> b) -> [a] -> [b]
map (Graph k cls color -> k -> Maybe (Node k cls color)
forall k cls color.
Uniquable k =>
Graph k cls color -> k -> Maybe (Node k cls color)
lookupNode Graph k cls color
graph)
                        ([k] -> [Maybe (Node k cls color)])
-> [k] -> [Maybe (Node k cls color)]
forall a b. (a -> b) -> a -> b
$ UniqSet k -> [k]
forall elt. UniqSet elt -> [elt]
nonDetEltsUniqSet
                        (UniqSet k -> [k]) -> UniqSet k -> [k]
forall a b. (a -> b) -> a -> b
$ Node k cls color -> UniqSet k
forall k cls color. Node k cls color -> UniqSet k
nodeConflicts Node k cls color
node
                        -- See Note [Unique Determinism and code generation]

        colors_conflict :: UniqSet color
colors_conflict = [color] -> UniqSet color
forall a. Uniquable a => [a] -> UniqSet a
mkUniqSet
                        ([color] -> UniqSet color) -> [color] -> UniqSet color
forall a b. (a -> b) -> a -> b
$ [Maybe color] -> [color]
forall a. [Maybe a] -> [a]
catMaybes
                        ([Maybe color] -> [color]) -> [Maybe color] -> [color]
forall a b. (a -> b) -> a -> b
$ (Node k cls color -> Maybe color)
-> [Node k cls color] -> [Maybe color]
forall a b. (a -> b) -> [a] -> [b]
map Node k cls color -> Maybe color
forall k cls color. Node k cls color -> Maybe color
nodeColor [Node k cls color]
nsConflicts

        -- the prefs of our neighbors
        colors_neighbor_prefs :: UniqSet color
colors_neighbor_prefs
                        = [color] -> UniqSet color
forall a. Uniquable a => [a] -> UniqSet a
mkUniqSet
                        ([color] -> UniqSet color) -> [color] -> UniqSet color
forall a b. (a -> b) -> a -> b
$ [[color]] -> [color]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[color]] -> [color]) -> [[color]] -> [color]
forall a b. (a -> b) -> a -> b
$ (Node k cls color -> [color]) -> [Node k cls color] -> [[color]]
forall a b. (a -> b) -> [a] -> [b]
map Node k cls color -> [color]
forall k cls color. Node k cls color -> [color]
nodePreference [Node k cls color]
nsConflicts

        -- colors that are still valid for us
        colors_ok_ex :: UniqSet color
colors_ok_ex    = UniqSet color -> UniqSet color -> UniqSet color
forall a. UniqSet a -> UniqSet a -> UniqSet a
minusUniqSet UniqSet color
colors_avail (Node k cls color -> UniqSet color
forall k cls color. Node k cls color -> UniqSet color
nodeExclusions Node k cls color
node)
        colors_ok :: UniqSet color
colors_ok       = UniqSet color -> UniqSet color -> UniqSet color
forall a. UniqSet a -> UniqSet a -> UniqSet a
minusUniqSet UniqSet color
colors_ok_ex UniqSet color
colors_conflict

        -- the colors that we prefer, and are still ok
        colors_ok_pref :: UniqSet color
colors_ok_pref  = UniqSet color -> UniqSet color -> UniqSet color
forall a. UniqSet a -> UniqSet a -> UniqSet a
intersectUniqSets
                                ([color] -> UniqSet color
forall a. Uniquable a => [a] -> UniqSet a
mkUniqSet ([color] -> UniqSet color) -> [color] -> UniqSet color
forall a b. (a -> b) -> a -> b
$ Node k cls color -> [color]
forall k cls color. Node k cls color -> [color]
nodePreference Node k cls color
node) UniqSet color
colors_ok

        -- the colors that we could choose while being nice to our neighbors
        colors_ok_nice :: UniqSet color
colors_ok_nice  = UniqSet color -> UniqSet color -> UniqSet color
forall a. UniqSet a -> UniqSet a -> UniqSet a
minusUniqSet
                                UniqSet color
colors_ok UniqSet color
colors_neighbor_prefs

        -- the best of all possible worlds..
        colors_ok_pref_nice :: UniqSet color
colors_ok_pref_nice
                        = UniqSet color -> UniqSet color -> UniqSet color
forall a. UniqSet a -> UniqSet a -> UniqSet a
intersectUniqSets
                                UniqSet color
colors_ok_nice UniqSet color
colors_ok_pref

        -- make the decision
        chooseColor :: Maybe color
chooseColor

                -- everyone is happy, yay!
                | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ UniqSet color -> Bool
forall a. UniqSet a -> Bool
isEmptyUniqSet UniqSet color
colors_ok_pref_nice
                , color
c : [color]
_         <- (color -> Bool) -> [color] -> [color]
forall a. (a -> Bool) -> [a] -> [a]
filter (\color
x -> color -> UniqSet color -> Bool
forall a. Uniquable a => a -> UniqSet a -> Bool
elementOfUniqSet color
x UniqSet color
colors_ok_pref_nice)
                                        (Node k cls color -> [color]
forall k cls color. Node k cls color -> [color]
nodePreference Node k cls color
node)
                = color -> Maybe color
forall a. a -> Maybe a
Just color
c

                -- we've got one of our preferences
                | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ UniqSet color -> Bool
forall a. UniqSet a -> Bool
isEmptyUniqSet UniqSet color
colors_ok_pref
                , color
c : [color]
_         <- (color -> Bool) -> [color] -> [color]
forall a. (a -> Bool) -> [a] -> [a]
filter (\color
x -> color -> UniqSet color -> Bool
forall a. Uniquable a => a -> UniqSet a -> Bool
elementOfUniqSet color
x UniqSet color
colors_ok_pref)
                                        (Node k cls color -> [color]
forall k cls color. Node k cls color -> [color]
nodePreference Node k cls color
node)
                = color -> Maybe color
forall a. a -> Maybe a
Just color
c

                -- it wasn't a preference, but it was still ok
                | Bool -> Bool
not (Bool -> Bool) -> Bool -> Bool
forall a b. (a -> b) -> a -> b
$ UniqSet color -> Bool
forall a. UniqSet a -> Bool
isEmptyUniqSet UniqSet color
colors_ok
                , color
c : [color]
_         <- UniqSet color -> [color]
forall elt. UniqSet elt -> [elt]
nonDetEltsUniqSet UniqSet color
colors_ok
                -- See Note [Unique Determinism and code generation]
                = color -> Maybe color
forall a. a -> Maybe a
Just color
c

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

   in   Maybe color
chooseColor