utils/GraphOps.hs

    -- | 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 =
              nonDetFoldUniqSet
                -- It's OK to use nonDetFoldUFM here because the
                -- operation is commutative
                (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 =
              nonDetFoldUniqSet
                -- It's OK to use nonDetFoldUFM here because the
                -- operation is commutative
                (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)
            => 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
                            $ nonDetEltsUniqSet (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
                            $ nonDetEltsUniqSet (nodeCoalesce node)
                            -- See Note [Unique Determinism and code generation]

                    -- 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    :: Graph k cls color -> Int

    size graph
            = sizeUFM $ graphMap graph

    -- | Union two graphs together.
    union   :: 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 : [])      <- nonDetEltsUniqSet conflicts
            = graphMapModify
            $ adjustWithDefaultUFM
                    id
                    (newNode u (getClass u))
                    u

            | otherwise
            = graphMapModify
            $ \fm -> foldl' (\g u  -> addConflictSet1 u getClass conflicts g) fm
                    $ nonDetEltsUniqSet conflicts
                    -- See Note [Unique Determinism and code generation]

    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 aggressive coalescing on this graph.
    --      returns the new graph and the list of pairs of nodes that got coalesced 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))
                    $ nonDetEltsUFM $ graphMap graph
                    -- See Note [Unique Determinism and code generation]

            -- 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    <- nonDetEltsUniqSet $ nodeCoalesce node1 ]
                            -- See Note [Unique Determinism and code generation]

            -- 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 \/ aggressively.
    --      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)
            => 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, Eq cls)
            => 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
            => 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     = nonDetFoldUniqSet (adjustUFM_C (freezeEdge k)) fm1
                        -- It's OK to use nonDetFoldUFM here because the operation
                        -- is commutative
                            $ 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)
            => 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)
            => Graph k cls color
            -> Graph k cls color

    freezeAllInGraph graph
            = foldr freezeNode graph
                    $ map nodeId
                    $ nonDetEltsUFM $ graphMap graph
                    -- See Note [Unique Determinism and code generation]

    -- | Find all the nodes in the graph that meet some criteria
    --
    scanGraph
            :: (Node k cls color -> Bool)
            -> Graph k cls color
            -> [Node k cls color]

    scanGraph match graph
            = filter match $ nonDetEltsUFM $ graphMap graph
              -- See Note [Unique Determinism and code generation]

    -- | 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       $ nonDetEltsUFM $ graphMap graph)
                                    ++ (map nodeCoalesce        $ nonDetEltsUFM $ graphMap graph))

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

            -- Check that no conflicting nodes have the same color
            | badNodes      <- filter (not . (checkNode graph))
                            $ nonDetEltsUFM $ graphMap graph
                               -- See Note [Unique Determinism and code generation]
            , 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)
                            $  nonDetEltsUFM $ 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)
                                    $  nonDetEltsUniqSet $ nodeConflicts node
                -- See Note [Unique Determinism and code generation]

            , 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
            :: 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)))
            $ nonDetEltsUFM
            -- See Note [Unique Determinism and code generation]
            $ 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)