move generic graph-colouring code into util

It is needed by cmm/StackColor, and hence is needed even when there is no
native code generator.

1 files changed, 581 insertions, 0 deletions

diff --git a/compiler/utils/GraphOps.hs b/compiler/utils/GraphOps.hs
new file mode 100644
index 0000000000..414abe4efc
--- /dev/null
+++ b/compiler/utils/GraphOps.hs
@@ -0,0 +1,581 @@
+-- | Basic operations on graphs.
+--
+--	TODO: refine coalescing crieteria
+
+{-# OPTIONS -fno-warn-missing-signatures #-}
+
+module GraphOps (
+	addNode, 	delNode,	getNode,	lookupNode,	modNode,
+	size,
+	union,
+	addConflict,	delConflict,	addConflicts,
+	addCoalesce,	delCoalesce,	
+	addExclusion,	
+	addPreference,
+	coalesceNodes,	coalesceGraph,
+	freezeNode,	freezeOneInGraph, freezeAllInGraph,
+	scanGraph,
+	setColor,
+	validateGraph,
+	slurpNodeConflictCount
+)
+where
+
+import GraphBase
+
+import Outputable
+import Unique
+import UniqSet
+import UniqFM
+
+import Data.List	hiding (union)
+import Data.Maybe
+
+-- | Lookup a node from the graph.
+lookupNode 
+	:: Uniquable k
+	=> Graph k cls color
+	-> k -> Maybe (Node  k cls color)
+
+lookupNode graph k	
+	= lookupUFM (graphMap graph) k
+
+
+-- | Get a node from the graph, throwing an error if it's not there
+getNode
+	:: Uniquable k
+	=> Graph k cls color
+	-> k -> Node k cls color
+
+getNode graph k
+ = case lookupUFM (graphMap graph) k of
+	Just node	-> node
+	Nothing		-> panic "ColorOps.getNode: not found" 
+
+
+-- | Add a node to the graph, linking up its edges
+addNode :: Uniquable k
+	=> k -> Node k cls color 
+	-> Graph k cls color -> Graph k cls color
+	
+addNode k node graph
+ = let	
+ 	-- add back conflict edges from other nodes to this one
+ 	map_conflict	
+		= foldUniqSet 
+			(adjustUFM (\n -> n { nodeConflicts = addOneToUniqSet (nodeConflicts n) k}))
+			(graphMap graph)
+			(nodeConflicts node)
+			
+	-- add back coalesce edges from other nodes to this one
+	map_coalesce
+		= foldUniqSet
+			(adjustUFM (\n -> n { nodeCoalesce = addOneToUniqSet (nodeCoalesce n) k}))
+			map_conflict
+			(nodeCoalesce node)
+	
+  in	graph
+  	{ graphMap	= addToUFM map_coalesce k node}
+		
+
+-- | Delete a node and all its edges from the graph.
+delNode :: (Uniquable k, Outputable k)
+	=> k -> Graph k cls color -> Maybe (Graph k cls color)
+
+delNode k graph
+	| Just node	<- lookupNode graph k
+	= let	-- delete conflict edges from other nodes to this one.
+		graph1	= foldl' (\g k1 -> let Just g' = delConflict k1 k g in g') graph
+			$ uniqSetToList (nodeConflicts node)
+	
+		-- delete coalesce edge from other nodes to this one.
+		graph2	= foldl' (\g k1 -> let Just g' = delCoalesce k1 k g in g') graph1
+			$ uniqSetToList (nodeCoalesce node)
+	
+		-- delete the node
+		graph3	= graphMapModify (\fm -> delFromUFM fm k) graph2
+	
+	  in	Just graph3
+		
+	| otherwise
+	= Nothing
+
+
+-- | Modify a node in the graph.
+--	returns Nothing if the node isn't present.
+--
+modNode :: Uniquable k
+	=> (Node k cls color -> Node k cls color) 
+	-> k -> Graph k cls color -> Maybe (Graph k cls color)
+
+modNode f k graph
+ = case lookupNode graph k of
+ 	Just Node{}
+	 -> Just
+	 $  graphMapModify
+		 (\fm	-> let	Just node	= lookupUFM fm k
+			   	node'		= f node
+			   in	addToUFM fm k node') 
+		graph
+
+	Nothing	-> Nothing
+
+
+-- | Get the size of the graph, O(n)
+size	:: Uniquable k 
+	=> Graph k cls color -> Int
+	
+size graph	
+	= sizeUFM $ graphMap graph
+	
+
+-- | Union two graphs together.
+union	:: Uniquable k
+	=> Graph k cls color -> Graph k cls color -> Graph k cls color
+	
+union	graph1 graph2
+	= Graph 
+	{ graphMap		= plusUFM (graphMap graph1) (graphMap graph2) }
+
+
+-- | Add a conflict between nodes to the graph, creating the nodes required.
+--	Conflicts are virtual regs which need to be colored differently.
+addConflict
+	:: Uniquable k
+	=> (k, cls) -> (k, cls) 
+	-> Graph k cls color -> Graph k cls color
+
+addConflict (u1, c1) (u2, c2)
+ = let	addNeighbor u c u'
+	 	= adjustWithDefaultUFM
+			(\node -> node { nodeConflicts = addOneToUniqSet (nodeConflicts node) u' })
+			(newNode u c)  { nodeConflicts = unitUniqSet u' }
+			u
+	
+   in	graphMapModify
+   	( addNeighbor u1 c1 u2 
+	. addNeighbor u2 c2 u1)
+
+ 
+-- | Delete a conflict edge. k1 -> k2
+--	returns Nothing if the node isn't in the graph
+delConflict 
+	:: Uniquable k
+	=> k -> k
+	-> Graph k cls color -> Maybe (Graph k cls color)
+	
+delConflict k1 k2
+	= modNode
+		(\node -> node { nodeConflicts = delOneFromUniqSet (nodeConflicts node) k2 })
+		k1
+
+
+-- | Add some conflicts to the graph, creating nodes if required.
+--	All the nodes in the set are taken to conflict with each other.
+addConflicts
+	:: Uniquable k
+	=> UniqSet k -> (k -> cls)
+	-> Graph k cls color -> Graph k cls color
+	
+addConflicts conflicts getClass
+
+	-- just a single node, but no conflicts, create the node anyway.
+	| (u : [])	<- uniqSetToList conflicts
+	= graphMapModify 
+	$ adjustWithDefaultUFM 
+		id
+		(newNode u (getClass u)) 
+		u
+
+	| otherwise
+	= graphMapModify
+	$ (\fm -> foldl' (\g u  -> addConflictSet1 u getClass conflicts g) fm
+		$ uniqSetToList conflicts)
+
+
+addConflictSet1 u getClass set 
+ = case delOneFromUniqSet set u of
+    set' -> adjustWithDefaultUFM
+		(\node -> node 			{ nodeConflicts = unionUniqSets set' (nodeConflicts node) } )
+		(newNode u (getClass u))	{ nodeConflicts = set' }
+		u
+
+
+-- | Add an exclusion to the graph, creating nodes if required.
+--	These are extra colors that the node cannot use.
+addExclusion
+	:: (Uniquable k, Uniquable color)
+	=> k -> (k -> cls) -> color 
+	-> Graph k cls color -> Graph k cls color
+	
+addExclusion u getClass color 
+ 	= graphMapModify
+	$ adjustWithDefaultUFM 
+		(\node -> node 			{ nodeExclusions = addOneToUniqSet (nodeExclusions node) color })
+		(newNode u (getClass u))  	{ nodeExclusions = unitUniqSet color }
+		u
+
+
+-- | Add a coalescence edge to the graph, creating nodes if requried.
+--	It is considered adventageous to assign the same color to nodes in a coalesence.
+addCoalesce 
+	:: Uniquable k
+	=> (k, cls) -> (k, cls) 
+	-> Graph k cls color -> Graph k cls color
+	
+addCoalesce (u1, c1) (u2, c2) 
+ = let	addCoalesce u c u'
+ 	 = 	adjustWithDefaultUFM
+	 		(\node -> node { nodeCoalesce = addOneToUniqSet (nodeCoalesce node) u' })
+			(newNode u c)  { nodeCoalesce = unitUniqSet u' }
+			u
+			
+   in	graphMapModify
+     	( addCoalesce u1 c1 u2
+        . addCoalesce u2 c2 u1)
+
+
+-- | Delete a coalescence edge (k1 -> k2) from the graph.
+delCoalesce
+	:: Uniquable k
+	=> k -> k 
+	-> Graph k cls color	-> Maybe (Graph k cls color)
+
+delCoalesce k1 k2
+	= modNode (\node -> node { nodeCoalesce = delOneFromUniqSet (nodeCoalesce node) k2 })
+		k1
+
+
+-- | Add a color preference to the graph, creating nodes if required.
+--	The most recently added preference is the most prefered.
+--	The algorithm tries to assign a node it's prefered color if possible.
+--
+addPreference 
+	:: Uniquable k
+	=> (k, cls) -> color
+	-> Graph k cls color -> Graph k cls color
+	
+addPreference (u, c) color 
+ 	= graphMapModify
+	$ adjustWithDefaultUFM 
+		(\node -> node { nodePreference = color : (nodePreference node) })
+		(newNode u c)  { nodePreference = [color] }
+		u
+
+
+-- | Do agressive coalescing on this graph.
+--	returns	the new graph and the list of pairs of nodes that got coaleced together.
+--	for each pair, the resulting node will have the least key and be second in the pair.
+--
+coalesceGraph
+	:: (Uniquable k, Ord k, Eq cls, Outputable k)
+	=> Bool			-- ^ If True, coalesce nodes even if this might make the graph
+				--	less colorable (aggressive coalescing)
+	-> Triv k cls color
+	-> Graph k cls color
+	-> (Graph k cls color, [(k, k)])
+
+coalesceGraph aggressive triv graph
+	= coalesceGraph' aggressive triv graph []
+
+coalesceGraph' aggressive triv graph kkPairsAcc
+ = let
+ 	-- find all the nodes that have coalescence edges
+	cNodes	= filter (\node -> not $ isEmptyUniqSet (nodeCoalesce node))
+		$ eltsUFM $ graphMap graph
+
+	-- build a list of pairs of keys for node's we'll try and coalesce
+	--	every pair of nodes will appear twice in this list
+	--	ie [(k1, k2), (k2, k1) ... ]
+	--	This is ok, GrapOps.coalesceNodes handles this and it's convenient for
+	--	build a list of what nodes get coalesced together for later on.
+	--
+	cList	= [ (nodeId node1, k2)
+			| node1	<- cNodes
+			, k2	<- uniqSetToList $ nodeCoalesce node1 ]
+
+	-- do the coalescing, returning the new graph and a list of pairs of keys
+	--	that got coalesced together.
+	(graph', mPairs)
+		= mapAccumL (coalesceNodes aggressive triv) graph cList
+
+	-- keep running until there are no more coalesces can be found
+   in	case catMaybes mPairs of
+	 []	-> (graph', kkPairsAcc)
+	 pairs	-> coalesceGraph' aggressive triv graph' (pairs ++ kkPairsAcc)
+
+
+-- | Coalesce this pair of nodes unconditionally / agressively.
+--	The resulting node is the one with the least key.
+--
+--	returns: Just 	 the pair of keys if the nodes were coalesced
+--			 the second element of the pair being the least one
+--
+--		 Nothing if either of the nodes weren't in the graph
+
+coalesceNodes
+	:: (Uniquable k, Ord k, Eq cls, Outputable k)
+	=> Bool			-- ^ If True, coalesce nodes even if this might make the graph
+				--	less colorable (aggressive coalescing)
+	-> Triv  k cls color
+	-> Graph k cls color
+	-> (k, k)		-- ^ keys of the nodes to be coalesced
+	-> (Graph k cls color, Maybe (k, k))
+
+coalesceNodes aggressive triv graph (k1, k2)
+	| (kMin, kMax)	<- if k1 < k2
+				then (k1, k2)
+				else (k2, k1)
+
+	-- the nodes being coalesced must be in the graph
+	, Just nMin	<- lookupNode graph kMin
+	, Just nMax	<- lookupNode graph kMax
+
+	-- can't coalesce conflicting modes
+	, not $ elementOfUniqSet kMin (nodeConflicts nMax)
+	, not $ elementOfUniqSet kMax (nodeConflicts nMin)
+
+	= coalesceNodes_merge aggressive triv graph kMin kMax nMin nMax
+
+	-- don't do the coalescing after all
+	| otherwise
+	= (graph, Nothing)
+
+coalesceNodes_merge aggressive triv graph kMin kMax nMin nMax
+
+	-- sanity checks
+	| nodeClass nMin /= nodeClass nMax
+	= error "GraphOps.coalesceNodes: can't coalesce nodes of different classes."
+
+	| not (isNothing (nodeColor nMin) && isNothing (nodeColor nMax))
+	= error "GraphOps.coalesceNodes: can't coalesce colored nodes."
+
+	| nodeId nMin == nodeId nMax
+	= error "GraphOps.coalesceNodes: can't coalesce the same node."
+
+	---
+	| otherwise
+	= let
+		-- the new node gets all the edges from its two components
+		node	=
+		 Node	{ nodeId		= kMin
+			, nodeClass		= nodeClass nMin
+			, nodeColor		= Nothing
+
+			-- nodes don't conflict with themselves..
+			, nodeConflicts
+				= (unionUniqSets (nodeConflicts nMin) (nodeConflicts nMax))
+					`delOneFromUniqSet` kMin
+					`delOneFromUniqSet` kMax
+
+			, nodeExclusions	= unionUniqSets (nodeExclusions nMin) (nodeExclusions nMax)
+			, nodePreference	= nodePreference nMin ++ nodePreference nMax
+
+			-- nodes don't coalesce with themselves..
+			, nodeCoalesce
+				= (unionUniqSets (nodeCoalesce nMin) (nodeCoalesce nMax))
+					`delOneFromUniqSet` kMin
+					`delOneFromUniqSet` kMax
+			}
+
+	  in	coalesceNodes_check aggressive triv graph kMin kMax node
+
+coalesceNodes_check aggressive triv graph kMin kMax node
+
+	-- Unless we're coalescing aggressively, if the result node is not trivially
+	--	colorable then don't do the coalescing.
+	| not aggressive
+	, not $ triv (nodeClass node) (nodeConflicts node) (nodeExclusions node)
+	= (graph, Nothing)
+
+	| otherwise
+	= let -- delete the old nodes from the graph and add the new one
+		Just graph1	= delNode kMax graph
+		Just graph2	= delNode kMin graph1
+		graph3		= addNode kMin node graph2
+
+	  in	(graph3, Just (kMax, kMin))
+
+
+-- | Freeze a node
+--	This is for the iterative coalescer.
+--	By freezing a node we give up on ever coalescing it.
+--	Move all its coalesce edges into the frozen set - and update
+--	back edges from other nodes.
+--
+freezeNode
+	:: Uniquable k
+	=> k			-- ^ key of the node to freeze
+	-> Graph k cls color	-- ^ the graph
+	-> Graph k cls color	-- ^ graph with that node frozen
+
+freezeNode k
+  = graphMapModify
+  $ \fm ->
+    let
+    	-- freeze all the edges in the node to be frozen
+    	Just node = lookupUFM fm k
+	node'	= node
+		{ nodeCoalesce		= emptyUniqSet }
+
+	fm1	= addToUFM fm k node'
+
+	-- update back edges pointing to this node
+	freezeEdge k node
+	 = if elementOfUniqSet k (nodeCoalesce node)
+	 	then node
+			{ nodeCoalesce		= delOneFromUniqSet (nodeCoalesce node) k }
+		else	panic "GraphOps.freezeNode: edge to freeze wasn't in the coalesce set"
+
+	fm2	= foldUniqSet (adjustUFM (freezeEdge k)) fm1
+	   		$ nodeCoalesce node
+
+    in	fm2
+
+
+-- | Freeze one node in the graph
+--	This if for the iterative coalescer.
+--	Look for a move related node of low degree and freeze it.
+--
+--	We probably don't need to scan the whole graph looking for the node of absolute
+--	lowest degree. Just sample the first few and choose the one with the lowest 
+--	degree out of those. Also, we don't make any distinction between conflicts of different
+--	classes.. this is just a heuristic, after all.
+--
+--	IDEA:	freezing a node might free it up for Simplify.. would be good to check for triv
+--		right here, and add it to a worklist if known triv/non-move nodes.
+--
+freezeOneInGraph
+	:: (Uniquable k, Outputable k)
+	=> Graph k cls color
+	-> ( Graph k cls color		-- the new graph
+	   , Bool )			-- whether we found a node to freeze
+
+freezeOneInGraph graph
+ = let	compareNodeDegree n1 n2
+		= compare (sizeUniqSet $ nodeConflicts n1) (sizeUniqSet $ nodeConflicts n2)
+
+	candidates
+		= sortBy compareNodeDegree
+		$ take 5	-- 5 isn't special, it's just a small number.
+		$ scanGraph (\node -> not $ isEmptyUniqSet (nodeCoalesce node)) graph
+
+   in	case candidates of
+
+	 -- there wasn't anything available to freeze
+   	 []	-> (graph, False)
+
+	 -- we found something to freeze
+	 (n : _)
+	  -> ( freezeNode (nodeId n) graph
+	     , True)
+
+
+-- | Freeze all the nodes in the graph
+--	for debugging the iterative allocator.
+--
+freezeAllInGraph
+	:: (Uniquable k, Outputable k)
+	=> Graph k cls color
+	-> Graph k cls color
+
+freezeAllInGraph graph
+	= foldr freezeNode graph
+   		$ map nodeId
+		$ eltsUFM $ graphMap graph
+
+
+-- | Find all the nodes in the graph that meet some criteria
+--
+scanGraph
+	:: Uniquable k
+	=> (Node k cls color -> Bool)
+	-> Graph k cls color
+	-> [Node k cls color]
+
+scanGraph match graph
+	= filter match $ eltsUFM $ graphMap graph
+
+
+-- | validate the internal structure of a graph
+--	all its edges should point to valid nodes
+--	if they don't then throw an error
+--
+validateGraph
+	:: (Uniquable k, Outputable k)
+	=> SDoc
+	-> Graph k cls color
+	-> Graph k cls color
+
+validateGraph doc graph
+ = let	edges	= unionManyUniqSets
+ 			(  (map nodeConflicts       $ eltsUFM $ graphMap graph)
+			++ (map nodeCoalesce        $ eltsUFM $ graphMap graph))
+
+	nodes	= mkUniqSet $ map nodeId $ eltsUFM $ graphMap graph
+	
+	badEdges = minusUniqSet edges nodes
+	
+  in	if isEmptyUniqSet badEdges 
+  	 then 	graph
+	 else	pprPanic "GraphOps.validateGraph"
+	 	( text  "-- bad edges"
+		$$ vcat (map ppr $ uniqSetToList badEdges)
+		$$ text "----------------------------"
+		$$ doc)
+
+
+-- | Slurp out a map of how many nodes had a certain number of conflict neighbours
+
+slurpNodeConflictCount
+	:: Uniquable k
+	=> Graph k cls color
+	-> UniqFM (Int, Int)	-- ^ (conflict neighbours, num nodes with that many conflicts)
+
+slurpNodeConflictCount graph
+	= addListToUFM_C
+		(\(c1, n1) (_, n2) -> (c1, n1 + n2))
+		emptyUFM
+	$ map 	(\node
+	 	  -> let count	= sizeUniqSet $ nodeConflicts node
+		     in  (count, (count, 1)))
+	$ eltsUFM
+	$ graphMap graph
+
+
+-- | Set the color of a certain node
+setColor 
+	:: Uniquable k
+	=> k -> color
+	-> Graph k cls color -> Graph k cls color
+	
+setColor u color
+ 	= graphMapModify
+ 	$ adjustUFM
+		(\n -> n { nodeColor = Just color })
+		u 
+	
+
+{-# INLINE 	adjustWithDefaultUFM #-}
+adjustWithDefaultUFM 
+	:: Uniquable k 
+	=> (a -> a) -> a -> k 
+	-> UniqFM a -> UniqFM a
+
+adjustWithDefaultUFM f def k map
+	= addToUFM_C 
+		(\old _ -> f old)
+		map
+		k def
+		
+{-# INLINE adjustUFM #-}
+adjustUFM 
+	:: Uniquable k
+	=> (a -> a)
+	-> k -> UniqFM a -> UniqFM a
+
+adjustUFM f k map
+ = case lookupUFM map k of
+ 	Nothing	-> map
+	Just a	-> addToUFM map k (f a)
+