Make some more modules use LazyUniqFM instead of UniqFM
[ghc-hetmet.git] / compiler / utils / GraphOps.hs
1 {-# OPTIONS -fno-warn-missing-signatures #-}
2 -- | Basic operations on graphs.
3 --
4
5 module GraphOps (
6         addNode,        delNode,        getNode,        lookupNode,     modNode,
7         size,
8         union,
9         addConflict,    delConflict,    addConflicts,
10         addCoalesce,    delCoalesce,    
11         addExclusion,   
12         addPreference,
13         coalesceNodes,  coalesceGraph,
14         freezeNode,     freezeOneInGraph, freezeAllInGraph,
15         scanGraph,
16         setColor,
17         validateGraph,
18         slurpNodeConflictCount
19 )
20 where
21
22 import GraphBase
23
24 import Outputable
25 import Unique
26 import UniqSet
27 import LazyUniqFM
28
29 import Data.List        hiding (union)
30 import Data.Maybe
31
32 -- | Lookup a node from the graph.
33 lookupNode 
34         :: Uniquable k
35         => Graph k cls color
36         -> k -> Maybe (Node  k cls color)
37
38 lookupNode graph k      
39         = lookupUFM (graphMap graph) k
40
41
42 -- | Get a node from the graph, throwing an error if it's not there
43 getNode
44         :: Uniquable k
45         => Graph k cls color
46         -> k -> Node k cls color
47
48 getNode graph k
49  = case lookupUFM (graphMap graph) k of
50         Just node       -> node
51         Nothing         -> panic "ColorOps.getNode: not found" 
52
53
54 -- | Add a node to the graph, linking up its edges
55 addNode :: Uniquable k
56         => k -> Node k cls color 
57         -> Graph k cls color -> Graph k cls color
58         
59 addNode k node graph
60  = let  
61         -- add back conflict edges from other nodes to this one
62         map_conflict    
63                 = foldUniqSet 
64                         (adjustUFM (\n -> n { nodeConflicts = addOneToUniqSet (nodeConflicts n) k}))
65                         (graphMap graph)
66                         (nodeConflicts node)
67                         
68         -- add back coalesce edges from other nodes to this one
69         map_coalesce
70                 = foldUniqSet
71                         (adjustUFM (\n -> n { nodeCoalesce = addOneToUniqSet (nodeCoalesce n) k}))
72                         map_conflict
73                         (nodeCoalesce node)
74         
75   in    graph
76         { graphMap      = addToUFM map_coalesce k node}
77                 
78
79 -- | Delete a node and all its edges from the graph.
80 delNode :: (Uniquable k, Outputable k)
81         => k -> Graph k cls color -> Maybe (Graph k cls color)
82
83 delNode k graph
84         | Just node     <- lookupNode graph k
85         = let   -- delete conflict edges from other nodes to this one.
86                 graph1  = foldl' (\g k1 -> let Just g' = delConflict k1 k g in g') graph
87                         $ uniqSetToList (nodeConflicts node)
88         
89                 -- delete coalesce edge from other nodes to this one.
90                 graph2  = foldl' (\g k1 -> let Just g' = delCoalesce k1 k g in g') graph1
91                         $ uniqSetToList (nodeCoalesce node)
92         
93                 -- delete the node
94                 graph3  = graphMapModify (\fm -> delFromUFM fm k) graph2
95         
96           in    Just graph3
97                 
98         | otherwise
99         = Nothing
100
101
102 -- | Modify a node in the graph.
103 --      returns Nothing if the node isn't present.
104 --
105 modNode :: Uniquable k
106         => (Node k cls color -> Node k cls color) 
107         -> k -> Graph k cls color -> Maybe (Graph k cls color)
108
109 modNode f k graph
110  = case lookupNode graph k of
111         Just Node{}
112          -> Just
113          $  graphMapModify
114                  (\fm   -> let  Just node       = lookupUFM fm k
115                                 node'           = f node
116                            in   addToUFM fm k node') 
117                 graph
118
119         Nothing -> Nothing
120
121
122 -- | Get the size of the graph, O(n)
123 size    :: Uniquable k 
124         => Graph k cls color -> Int
125         
126 size graph      
127         = sizeUFM $ graphMap graph
128         
129
130 -- | Union two graphs together.
131 union   :: Uniquable k
132         => Graph k cls color -> Graph k cls color -> Graph k cls color
133         
134 union   graph1 graph2
135         = Graph 
136         { graphMap              = plusUFM (graphMap graph1) (graphMap graph2) }
137
138
139 -- | Add a conflict between nodes to the graph, creating the nodes required.
140 --      Conflicts are virtual regs which need to be colored differently.
141 addConflict
142         :: Uniquable k
143         => (k, cls) -> (k, cls) 
144         -> Graph k cls color -> Graph k cls color
145
146 addConflict (u1, c1) (u2, c2)
147  = let  addNeighbor u c u'
148                 = adjustWithDefaultUFM
149                         (\node -> node { nodeConflicts = addOneToUniqSet (nodeConflicts node) u' })
150                         (newNode u c)  { nodeConflicts = unitUniqSet u' }
151                         u
152         
153    in   graphMapModify
154         ( addNeighbor u1 c1 u2 
155         . addNeighbor u2 c2 u1)
156
157  
158 -- | Delete a conflict edge. k1 -> k2
159 --      returns Nothing if the node isn't in the graph
160 delConflict 
161         :: Uniquable k
162         => k -> k
163         -> Graph k cls color -> Maybe (Graph k cls color)
164         
165 delConflict k1 k2
166         = modNode
167                 (\node -> node { nodeConflicts = delOneFromUniqSet (nodeConflicts node) k2 })
168                 k1
169
170
171 -- | Add some conflicts to the graph, creating nodes if required.
172 --      All the nodes in the set are taken to conflict with each other.
173 addConflicts
174         :: Uniquable k
175         => UniqSet k -> (k -> cls)
176         -> Graph k cls color -> Graph k cls color
177         
178 addConflicts conflicts getClass
179
180         -- just a single node, but no conflicts, create the node anyway.
181         | (u : [])      <- uniqSetToList conflicts
182         = graphMapModify 
183         $ adjustWithDefaultUFM 
184                 id
185                 (newNode u (getClass u)) 
186                 u
187
188         | otherwise
189         = graphMapModify
190         $ (\fm -> foldl' (\g u  -> addConflictSet1 u getClass conflicts g) fm
191                 $ uniqSetToList conflicts)
192
193
194 addConflictSet1 u getClass set 
195  = case delOneFromUniqSet set u of
196     set' -> adjustWithDefaultUFM
197                 (\node -> node                  { nodeConflicts = unionUniqSets set' (nodeConflicts node) } )
198                 (newNode u (getClass u))        { nodeConflicts = set' }
199                 u
200
201
202 -- | Add an exclusion to the graph, creating nodes if required.
203 --      These are extra colors that the node cannot use.
204 addExclusion
205         :: (Uniquable k, Uniquable color)
206         => k -> (k -> cls) -> color 
207         -> Graph k cls color -> Graph k cls color
208         
209 addExclusion u getClass color 
210         = graphMapModify
211         $ adjustWithDefaultUFM 
212                 (\node -> node                  { nodeExclusions = addOneToUniqSet (nodeExclusions node) color })
213                 (newNode u (getClass u))        { nodeExclusions = unitUniqSet color }
214                 u
215
216
217 -- | Add a coalescence edge to the graph, creating nodes if requried.
218 --      It is considered adventageous to assign the same color to nodes in a coalesence.
219 addCoalesce 
220         :: Uniquable k
221         => (k, cls) -> (k, cls) 
222         -> Graph k cls color -> Graph k cls color
223         
224 addCoalesce (u1, c1) (u2, c2) 
225  = let  addCoalesce u c u'
226          =      adjustWithDefaultUFM
227                         (\node -> node { nodeCoalesce = addOneToUniqSet (nodeCoalesce node) u' })
228                         (newNode u c)  { nodeCoalesce = unitUniqSet u' }
229                         u
230                         
231    in   graphMapModify
232         ( addCoalesce u1 c1 u2
233         . addCoalesce u2 c2 u1)
234
235
236 -- | Delete a coalescence edge (k1 -> k2) from the graph.
237 delCoalesce
238         :: Uniquable k
239         => k -> k 
240         -> Graph k cls color    -> Maybe (Graph k cls color)
241
242 delCoalesce k1 k2
243         = modNode (\node -> node { nodeCoalesce = delOneFromUniqSet (nodeCoalesce node) k2 })
244                 k1
245
246
247 -- | Add a color preference to the graph, creating nodes if required.
248 --      The most recently added preference is the most prefered.
249 --      The algorithm tries to assign a node it's prefered color if possible.
250 --
251 addPreference 
252         :: Uniquable k
253         => (k, cls) -> color
254         -> Graph k cls color -> Graph k cls color
255         
256 addPreference (u, c) color 
257         = graphMapModify
258         $ adjustWithDefaultUFM 
259                 (\node -> node { nodePreference = color : (nodePreference node) })
260                 (newNode u c)  { nodePreference = [color] }
261                 u
262
263
264 -- | Do agressive coalescing on this graph.
265 --      returns the new graph and the list of pairs of nodes that got coaleced together.
266 --      for each pair, the resulting node will have the least key and be second in the pair.
267 --
268 coalesceGraph
269         :: (Uniquable k, Ord k, Eq cls, Outputable k)
270         => Bool                 -- ^ If True, coalesce nodes even if this might make the graph
271                                 --      less colorable (aggressive coalescing)
272         -> Triv k cls color
273         -> Graph k cls color
274         -> ( Graph k cls color
275            , [(k, k)])          -- pairs of nodes that were coalesced, in the order that the
276                                 --      coalescing was applied.
277
278 coalesceGraph aggressive triv graph
279         = coalesceGraph' aggressive triv graph []
280
281 coalesceGraph' aggressive triv graph kkPairsAcc
282  = let
283         -- find all the nodes that have coalescence edges
284         cNodes  = filter (\node -> not $ isEmptyUniqSet (nodeCoalesce node))
285                 $ eltsUFM $ graphMap graph
286
287         -- build a list of pairs of keys for node's we'll try and coalesce
288         --      every pair of nodes will appear twice in this list
289         --      ie [(k1, k2), (k2, k1) ... ]
290         --      This is ok, GrapOps.coalesceNodes handles this and it's convenient for
291         --      build a list of what nodes get coalesced together for later on.
292         --
293         cList   = [ (nodeId node1, k2)
294                         | node1 <- cNodes
295                         , k2    <- uniqSetToList $ nodeCoalesce node1 ]
296
297         -- do the coalescing, returning the new graph and a list of pairs of keys
298         --      that got coalesced together.
299         (graph', mPairs)
300                 = mapAccumL (coalesceNodes aggressive triv) graph cList
301
302         -- keep running until there are no more coalesces can be found
303    in   case catMaybes mPairs of
304          []     -> (graph', reverse kkPairsAcc)
305          pairs  -> coalesceGraph' aggressive triv graph' (reverse pairs ++ kkPairsAcc)
306
307
308 -- | Coalesce this pair of nodes unconditionally / agressively.
309 --      The resulting node is the one with the least key.
310 --
311 --      returns: Just    the pair of keys if the nodes were coalesced
312 --                       the second element of the pair being the least one
313 --
314 --               Nothing if either of the nodes weren't in the graph
315
316 coalesceNodes
317         :: (Uniquable k, Ord k, Eq cls, Outputable k)
318         => Bool                 -- ^ If True, coalesce nodes even if this might make the graph
319                                 --      less colorable (aggressive coalescing)
320         -> Triv  k cls color
321         -> Graph k cls color
322         -> (k, k)               -- ^ keys of the nodes to be coalesced
323         -> (Graph k cls color, Maybe (k, k))
324
325 coalesceNodes aggressive triv graph (k1, k2)
326         | (kMin, kMax)  <- if k1 < k2
327                                 then (k1, k2)
328                                 else (k2, k1)
329
330         -- the nodes being coalesced must be in the graph
331         , Just nMin     <- lookupNode graph kMin
332         , Just nMax     <- lookupNode graph kMax
333
334         -- can't coalesce conflicting modes
335         , not $ elementOfUniqSet kMin (nodeConflicts nMax)
336         , not $ elementOfUniqSet kMax (nodeConflicts nMin)
337
338         = coalesceNodes_merge aggressive triv graph kMin kMax nMin nMax
339
340         -- don't do the coalescing after all
341         | otherwise
342         = (graph, Nothing)
343
344 coalesceNodes_merge aggressive triv graph kMin kMax nMin nMax
345
346         -- sanity checks
347         | nodeClass nMin /= nodeClass nMax
348         = error "GraphOps.coalesceNodes: can't coalesce nodes of different classes."
349
350         | not (isNothing (nodeColor nMin) && isNothing (nodeColor nMax))
351         = error "GraphOps.coalesceNodes: can't coalesce colored nodes."
352
353         | nodeId nMin == nodeId nMax
354         = error "GraphOps.coalesceNodes: can't coalesce the same node."
355
356         ---
357         | otherwise
358         = let
359                 -- the new node gets all the edges from its two components
360                 node    =
361                  Node   { nodeId                = kMin
362                         , nodeClass             = nodeClass nMin
363                         , nodeColor             = Nothing
364
365                         -- nodes don't conflict with themselves..
366                         , nodeConflicts
367                                 = (unionUniqSets (nodeConflicts nMin) (nodeConflicts nMax))
368                                         `delOneFromUniqSet` kMin
369                                         `delOneFromUniqSet` kMax
370
371                         , nodeExclusions        = unionUniqSets (nodeExclusions nMin) (nodeExclusions nMax)
372                         , nodePreference        = nodePreference nMin ++ nodePreference nMax
373
374                         -- nodes don't coalesce with themselves..
375                         , nodeCoalesce
376                                 = (unionUniqSets (nodeCoalesce nMin) (nodeCoalesce nMax))
377                                         `delOneFromUniqSet` kMin
378                                         `delOneFromUniqSet` kMax
379                         }
380
381           in    coalesceNodes_check aggressive triv graph kMin kMax node
382
383 coalesceNodes_check aggressive triv graph kMin kMax node
384
385         -- Unless we're coalescing aggressively, if the result node is not trivially
386         --      colorable then don't do the coalescing.
387         | not aggressive
388         , not $ triv (nodeClass node) (nodeConflicts node) (nodeExclusions node)
389         = (graph, Nothing)
390
391         | otherwise
392         = let -- delete the old nodes from the graph and add the new one
393                 Just graph1     = delNode kMax graph
394                 Just graph2     = delNode kMin graph1
395                 graph3          = addNode kMin node graph2
396
397           in    (graph3, Just (kMax, kMin))
398
399
400 -- | Freeze a node
401 --      This is for the iterative coalescer.
402 --      By freezing a node we give up on ever coalescing it.
403 --      Move all its coalesce edges into the frozen set - and update
404 --      back edges from other nodes.
405 --
406 freezeNode
407         :: Uniquable k
408         => k                    -- ^ key of the node to freeze
409         -> Graph k cls color    -- ^ the graph
410         -> Graph k cls color    -- ^ graph with that node frozen
411
412 freezeNode k
413   = graphMapModify
414   $ \fm ->
415     let
416         -- freeze all the edges in the node to be frozen
417         Just node = lookupUFM fm k
418         node'   = node
419                 { nodeCoalesce          = emptyUniqSet }
420
421         fm1     = addToUFM fm k node'
422
423         -- update back edges pointing to this node
424         freezeEdge k node
425          = if elementOfUniqSet k (nodeCoalesce node)
426                 then node
427                         { nodeCoalesce          = delOneFromUniqSet (nodeCoalesce node) k }
428                 else    panic "GraphOps.freezeNode: edge to freeze wasn't in the coalesce set"
429
430         fm2     = foldUniqSet (adjustUFM (freezeEdge k)) fm1
431                         $ nodeCoalesce node
432
433     in  fm2
434
435
436 -- | Freeze one node in the graph
437 --      This if for the iterative coalescer.
438 --      Look for a move related node of low degree and freeze it.
439 --
440 --      We probably don't need to scan the whole graph looking for the node of absolute
441 --      lowest degree. Just sample the first few and choose the one with the lowest 
442 --      degree out of those. Also, we don't make any distinction between conflicts of different
443 --      classes.. this is just a heuristic, after all.
444 --
445 --      IDEA:   freezing a node might free it up for Simplify.. would be good to check for triv
446 --              right here, and add it to a worklist if known triv/non-move nodes.
447 --
448 freezeOneInGraph
449         :: (Uniquable k, Outputable k)
450         => Graph k cls color
451         -> ( Graph k cls color          -- the new graph
452            , Bool )                     -- whether we found a node to freeze
453
454 freezeOneInGraph graph
455  = let  compareNodeDegree n1 n2
456                 = compare (sizeUniqSet $ nodeConflicts n1) (sizeUniqSet $ nodeConflicts n2)
457
458         candidates
459                 = sortBy compareNodeDegree
460                 $ take 5        -- 5 isn't special, it's just a small number.
461                 $ scanGraph (\node -> not $ isEmptyUniqSet (nodeCoalesce node)) graph
462
463    in   case candidates of
464
465          -- there wasn't anything available to freeze
466          []     -> (graph, False)
467
468          -- we found something to freeze
469          (n : _)
470           -> ( freezeNode (nodeId n) graph
471              , True)
472
473
474 -- | Freeze all the nodes in the graph
475 --      for debugging the iterative allocator.
476 --
477 freezeAllInGraph
478         :: (Uniquable k, Outputable k)
479         => Graph k cls color
480         -> Graph k cls color
481
482 freezeAllInGraph graph
483         = foldr freezeNode graph
484                 $ map nodeId
485                 $ eltsUFM $ graphMap graph
486
487
488 -- | Find all the nodes in the graph that meet some criteria
489 --
490 scanGraph
491         :: Uniquable k
492         => (Node k cls color -> Bool)
493         -> Graph k cls color
494         -> [Node k cls color]
495
496 scanGraph match graph
497         = filter match $ eltsUFM $ graphMap graph
498
499
500 -- | validate the internal structure of a graph
501 --      all its edges should point to valid nodes
502 --      If they don't then throw an error
503 --
504 validateGraph
505         :: (Uniquable k, Outputable k, Eq color)
506         => SDoc                         -- ^ extra debugging info to display on error
507         -> Bool                         -- ^ whether this graph is supposed to be colored.
508         -> Graph k cls color            -- ^ graph to validate
509         -> Graph k cls color            -- ^ validated graph
510
511 validateGraph doc isColored graph
512
513         -- Check that all edges point to valid nodes.
514         | edges         <- unionManyUniqSets
515                                 (  (map nodeConflicts       $ eltsUFM $ graphMap graph)
516                                 ++ (map nodeCoalesce        $ eltsUFM $ graphMap graph))
517
518         , nodes         <- mkUniqSet $ map nodeId $ eltsUFM $ graphMap graph
519         , badEdges      <- minusUniqSet edges nodes
520         , not $ isEmptyUniqSet badEdges
521         = pprPanic "GraphOps.validateGraph"
522                 (  text "Graph has edges that point to non-existant nodes"
523                 $$ text "  bad edges: " <> vcat (map ppr $ uniqSetToList badEdges)
524                 $$ doc )
525
526         -- Check that no conflicting nodes have the same color
527         | badNodes      <- filter (not . (checkNode graph))
528                         $ eltsUFM $ graphMap graph
529         , not $ null badNodes
530         = pprPanic "GraphOps.validateGraph"
531                 (  text "Node has same color as one of it's conflicts"
532                 $$ text "  bad nodes: " <> hcat (map (ppr . nodeId) badNodes)
533                 $$ doc)
534
535         -- If this is supposed to be a colored graph,
536         --      check that all nodes have a color.
537         | isColored
538         , badNodes      <- filter (\n -> isNothing $ nodeColor n)
539                         $  eltsUFM $ graphMap graph
540         , not $ null badNodes
541         = pprPanic "GraphOps.validateGraph"
542                 (  text "Supposably colored graph has uncolored nodes."
543                 $$ text "  uncolored nodes: " <> hcat (map (ppr . nodeId) badNodes)
544                 $$ doc )
545
546
547         -- graph looks ok
548         | otherwise
549         = graph
550
551
552 -- | If this node is colored, check that all the nodes which
553 --      conflict with it have different colors.
554 checkNode
555         :: (Uniquable k, Eq color)
556         => Graph k cls color
557         -> Node  k cls color
558         -> Bool                 -- ^ True if this node is ok
559         
560 checkNode graph node
561         | Just color            <- nodeColor node
562         , Just neighbors        <- sequence $ map (lookupNode graph)
563                                 $  uniqSetToList $ nodeConflicts node
564
565         , neighbourColors       <- catMaybes $ map nodeColor neighbors
566         , elem color neighbourColors
567         = False
568         
569         | otherwise
570         = True
571
572
573
574 -- | Slurp out a map of how many nodes had a certain number of conflict neighbours
575
576 slurpNodeConflictCount
577         :: Uniquable k
578         => Graph k cls color
579         -> UniqFM (Int, Int)    -- ^ (conflict neighbours, num nodes with that many conflicts)
580
581 slurpNodeConflictCount graph
582         = addListToUFM_C
583                 (\(c1, n1) (_, n2) -> (c1, n1 + n2))
584                 emptyUFM
585         $ map   (\node
586                   -> let count  = sizeUniqSet $ nodeConflicts node
587                      in  (count, (count, 1)))
588         $ eltsUFM
589         $ graphMap graph
590
591
592 -- | Set the color of a certain node
593 setColor 
594         :: Uniquable k
595         => k -> color
596         -> Graph k cls color -> Graph k cls color
597         
598 setColor u color
599         = graphMapModify
600         $ adjustUFM
601                 (\n -> n { nodeColor = Just color })
602                 u 
603         
604
605 {-# INLINE      adjustWithDefaultUFM #-}
606 adjustWithDefaultUFM 
607         :: Uniquable k 
608         => (a -> a) -> a -> k 
609         -> UniqFM a -> UniqFM a
610
611 adjustWithDefaultUFM f def k map
612         = addToUFM_C 
613                 (\old _ -> f old)
614                 map
615                 k def
616                 
617 {-# INLINE adjustUFM #-}
618 adjustUFM 
619         :: Uniquable k
620         => (a -> a)
621         -> k -> UniqFM a -> UniqFM a
622
623 adjustUFM f k map
624  = case lookupUFM map k of
625         Nothing -> map
626         Just a  -> addToUFM map k (f a)
627