{-# OPTIONS -fno-warn-missing-signatures #-}
-- | Basic operations on graphs.
--

module GraphOps (
	addNode, 	delNode,	getNode,	lookupNode,	modNode,
	size,
	union,
	addConflict,	delConflict,	addConflicts,
	addCoalesce,	delCoalesce,	
	addExclusion,	addExclusions,
	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_C (\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_C (\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

addExclusions
	:: (Uniquable k, Uniquable color)
	=> k -> (k -> cls) -> [color]
	-> Graph k cls color -> Graph k cls color

addExclusions u getClass colors graph
	= foldr (addExclusion u getClass) graph colors


-- | 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)])		-- pairs of nodes that were coalesced, in the order that the
	   			--	coalescing was applied.

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', reverse kkPairsAcc)
	 pairs	-> coalesceGraph' aggressive triv graph' (reverse 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)

	-- can't coalesce the same node
	, nodeId nMin /= nodeId nMax

	= 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."

	---
	| 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 node 	-- panic "GraphOps.freezeNode: edge to freeze wasn't in the coalesce set"
				-- If the edge isn't actually in the coelesce set then just ignore it.

	fm2	= foldUniqSet (adjustUFM_C (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, Eq color)
	=> SDoc				-- ^ extra debugging info to display on error
	-> Bool				-- ^ whether this graph is supposed to be colored.
	-> Graph k cls color		-- ^ graph to validate
	-> Graph k cls color		-- ^ validated graph

validateGraph doc isColored graph

	-- Check that all edges point to valid nodes.
  	| edges		<- unionManyUniqSets
 				(  (map nodeConflicts       $ eltsUFM $ graphMap graph)
				++ (map nodeCoalesce        $ eltsUFM $ graphMap graph))

	, nodes		<- mkUniqSet $ map nodeId $ eltsUFM $ graphMap graph
	, badEdges 	<- minusUniqSet edges nodes
	, not $ isEmptyUniqSet badEdges
	= pprPanic "GraphOps.validateGraph"
		(  text "Graph has edges that point to non-existant nodes"
		$$ text "  bad edges: " <> vcat (map ppr $ uniqSetToList badEdges)
		$$ doc )

	-- Check that no conflicting nodes have the same color
	| badNodes	<- filter (not . (checkNode graph))
			$ eltsUFM $ graphMap graph
	, not $ null badNodes
	= pprPanic "GraphOps.validateGraph"
		(  text "Node has same color as one of it's conflicts"
		$$ text "  bad nodes: " <> hcat (map (ppr . nodeId) badNodes)
		$$ doc)

	-- If this is supposed to be a colored graph,
	--	check that all nodes have a color.
	| isColored
	, badNodes	<- filter (\n -> isNothing $ nodeColor n)
			$  eltsUFM $ graphMap graph
	, not $ null badNodes
	= pprPanic "GraphOps.validateGraph"
		(  text "Supposably colored graph has uncolored nodes."
		$$ text "  uncolored nodes: " <> hcat (map (ppr . nodeId) badNodes)
		$$ doc )


	-- graph looks ok
	| otherwise
	= graph


-- | If this node is colored, check that all the nodes which
--	conflict with it have different colors.
checkNode
	:: (Uniquable k, Eq color)
	=> Graph k cls color
	-> Node  k cls color
	-> Bool			-- ^ True if this node is ok
	
checkNode graph node
	| Just color		<- nodeColor node
	, Just neighbors	<- sequence $ map (lookupNode graph)
				$  uniqSetToList $ nodeConflicts node

	, neighbourColors	<- catMaybes $ map nodeColor neighbors
	, elem color neighbourColors
	= False
	
	| otherwise
	= True



-- | 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_C
		(\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
		
-- Argument order different from UniqFM's adjustUFM
{-# INLINE adjustUFM_C #-}
adjustUFM_C 
	:: Uniquable k
	=> (a -> a)
	-> k -> UniqFM a -> UniqFM a

adjustUFM_C f k map
 = case lookupUFM map k of
 	Nothing	-> map
	Just a	-> addToUFM map k (f a)