    {-# OPTIONS -fno-warn-missing-signatures #-}

    -- | 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 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
    	-> 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 coaleced
    	   				--	 r1 should be replaced by r2 in the source

    colorGraph iterative spinCount colors triv spill 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_coalesced, kksCoalesce1)
    	 = if iterative
    		then (graph0, [])
    		else if spinCount == 0
    			then coalesceGraph True  triv graph0
    			else coalesceGraph False triv 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
    			$$ text "colors    = " <> ppr colors
    			$$ empty
    			$$ dotGraph (\_ -> text "white") triv graph_triv)

    	 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 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 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, reverse 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
    	, graph2	<- foldr (\k g -> let Just g' = delNode k g
    	   				  in  g')
    				graph ksTrivFound

    	= colorScan_spin iterative triv spill graph2
    		(ksTrivFound ++ ksTriv)
    		ksSpill
    		kksCoalesce

    	-- Coalesce:
    	-- 	If we're doing iterative coalescing and no triv nodes are avaliable
    	--	then it's time for a coalescing pass.
    	| iterative
    	= case coalesceGraph False triv graph of

    		-- we were able to coalesce something
    		--	go back to Simplify and see if this frees up more nodes to be trivially colorable.
    		(graph2, kksCoalesceFound @(_:_))
    		 -> colorScan_spin iterative triv spill graph2
    			ksTriv ksSpill (reverse 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, 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 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, 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 colors graph u
     = let	-- lookup the node
     	Just node	= lookupNode graph u

    	-- lookup the available colors for the class of this node.
    	colors_avail
    	 = case lookupUFM colors (nodeClass node) of
    	 	Nothing	-> pprPanic "selectColor: no colors available for class " (ppr (nodeClass node))
    		Just cs	-> cs

    	-- 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