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