-- | 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 :: Uniquable k => k -> (k -> cls) -> UniqSet k -> UniqFM (Node k cls color) -> UniqFM (Node k cls color) 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' :: (Uniquable k, Ord k, Eq cls, Outputable k) => Bool -> Triv k cls color -> Graph k cls color -> [(k, k)] -> ( Graph k cls color , [(k, k)]) 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 :: (Uniquable k, Ord k, Eq cls, Outputable k) => Bool -> Triv k cls color -> Graph k cls color -> k -> k -> Node k cls color -> Node k cls color -> (Graph k cls color, Maybe (k, k)) 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 :: (Uniquable k, Ord k, Eq cls, Outputable k) => Bool -> Triv k cls color -> Graph k cls color -> k -> k -> Node k cls color -> (Graph k cls color, Maybe (k, k)) 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)