content: -- | Basic operations on graphs.
  --
  
  {-# OPTIONS_GHC -Wno-incomplete-uni-patterns #-}
  
  module GHC.Data.Graph.Ops
     ( 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 GHC.Prelude
  
  import GHC.Data.Graph.Base
  
  import GHC.Utils.Outputable
  import GHC.Utils.Panic
  import GHC.Types.Unique
  import GHC.Types.Unique.Set
  import GHC.Types.Unique.FM
  
  import Data.List (mapAccumL, sortBy)
  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 =
            nonDetStrictFoldUniqSet
              -- It's OK to use a non-deterministic fold 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 =
            nonDetStrictFoldUniqSet
              -- It's OK to use a non-deterministic fold 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 k (Node k cls color)
                  -> UniqFM k (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 required.
  --      It is considered advantageous to assign the same color to nodes in a coalescence.
  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 preferred.
  --      The algorithm tries to assign a node it's preferred 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 "GHC.Data.Graph.Ops.coalesceNodes: can't coalesce nodes of different classes."
  
          | not (isNothing (nodeColor nMin) && isNothing (nodeColor nMax))
          = error "GHC.Data.Graph.Ops.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 "GHC.Data.Graph.Ops.freezeNode: edge to freeze wasn't in the coalesce set"
                                  -- If the edge isn't actually in the coalesce set then just ignore it.
  
          fm2     = nonDetStrictFoldUniqSet (adjustUFM_C (freezeEdge k)) fm1
                      -- It's OK to use a non-deterministic fold 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 "GHC.Data.Graph.Ops.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 "GHC.Data.Graph.Ops.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 "GHC.Data.Graph.Ops.validateGraph"
                  (  text "Supposedly 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       <- mapMaybe 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, 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 k a -> UniqFM k 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 k a -> UniqFM k a
  
  adjustUFM_C f k map
   = case lookupUFM map k of
          Nothing -> map
          Just a  -> addToUFM map k (f a)
  
  