--
-- Copyright (c) 2018 Andreas Klebinger
--

{-# LANGUAGE TypeFamilies, ScopedTypeVariables, CPP #-}

{-# OPTIONS_GHC -fprof-auto #-}
--{-# OPTIONS_GHC -ddump-simpl -ddump-to-file -ddump-cmm #-}

module BlockLayout
    ( sequenceTop )
where

#include "HsVersions.h"
import GhcPrelude

import Instruction
import NCGMonad
import CFG

import BlockId
import Cmm
import Hoopl.Collections
import Hoopl.Label
import Hoopl.Block

import DynFlags (gopt, GeneralFlag(..), DynFlags, backendMaintainsCfg)
import UniqFM
import Util
import Unique

import Digraph
import Outputable
import Maybes

-- DEBUGGING ONLY
--import Debug
--import Debug.Trace
import ListSetOps (removeDups)
import PprCmm ()

import OrdList
import Data.List
import Data.Foldable (toList)
import Hoopl.Graph

import qualified Data.Set as Set

{-
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  ~~~ Note [Chain based CFG serialization]
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  For additional information also look at
  https://ghc.haskell.org/trac/ghc/wiki/Commentary/Compiler/CodeLayout

  We have a CFG with edge weights based on which we try to place blocks next to
  each other.

  Edge weights not only represent likelyhood of control transfer between blocks
  but also how much a block would benefit from being placed sequentially after
  it's predecessor.
  For example blocks which are preceeded by an info table are more likely to end
  up in a different cache line than their predecessor. So there is less benefit
  in placing them sequentially.

  For example consider this example:

  A:  ...
      jmp cond D (weak successor)
      jmp B
  B:  ...
      jmp C
  C:  ...
      jmp X
  D:  ...
      jmp B (weak successor)

  We determine a block layout by building up chunks (calling them chains) of
  possible control flows for which blocks will be placed sequentially.

  Eg for our example we might end up with two chains like:
  [A->B->C->X],[D]. Blocks inside chains will always be placed sequentially.
  However there is no particular order in which chains are placed since
  (hopefully) the blocks for which sequentially is important have already
  been placed in the same chain.

  -----------------------------------------------------------------------------
      First try to create a lists of good chains.
  -----------------------------------------------------------------------------

  We do so by taking a block not yet placed in a chain and
  looking at these cases:

  *)  Check if the best predecessor of the block is at the end of a chain.
      If so add the current block to the end of that chain.

      Eg if we look at block C and already have the chain (A -> B)
      then we extend the chain to (A -> B -> C).

      Combined with the fact that we process blocks in reverse post order
      this means loop bodies and trivially sequential control flow already
      ends up as a single chain.

  *)  Otherwise we create a singleton chain from the block we are looking at.
      Eg if we have from the example above already constructed (A->B)
      and look at D we create the chain (D) resulting in the chains [A->B, D]

  -----------------------------------------------------------------------------
      We then try to fuse chains.
  -----------------------------------------------------------------------------

  There are edge cases which result in two chains being created which trivially
  represent linear control flow. For example we might have the chains
  [(A-B-C),(D-E)] with an cfg triangle:

      A----->C->D->E
       \->B-/

  We also get three independent chains if two branches end with a jump
  to a common successor.

  We take care of these cases by fusing chains which are connected by an
  edge.

  We do so by looking at the list of edges sorted by weight.
  Given the edge (C -> D) we try to find two chains such that:
      * C is at the end of chain one.
      * D is in front of chain two.
      * If two such chains exist we fuse them.
  We then remove the edge and repeat the process for the rest of the edges.

  -----------------------------------------------------------------------------
      Place indirect successors (neighbours) after each other
  -----------------------------------------------------------------------------

  We might have chains [A,B,C,X],[E] in a CFG of the sort:

    A ---> B ---> C --------> X(exit)
                   \- ->E- -/

  While E does not follow X it's still beneficial to place them near each other.
  This can be advantageous if eg C,X,E will end up in the same cache line.

  TODO: If we remove edges as we use them (eg if we build up A->B remove A->B
        from the list) we could save some more work in later phases.

  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  ~~~ Note [Triangle Control Flow]
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  Checking if an argument is already evaluating leads to a somewhat
  special case  which looks like this:

    A:
        if (R1 & 7 != 0) goto Leval; else goto Lwork;
    Leval: // global
        call (I64[R1])(R1) returns to Lwork, args: 8, res: 8, upd: 8;
    Lwork: // global
        ...

        A
        |\
        | Leval
        |/ - (This edge can be missing because of optimizations)
        Lwork

  Once we hit the metal the call instruction is just 2-3 bytes large
  depending on the register used. So we lay out the assembly like this:

        movq %rbx,%rax
        andl $7,%eax
        cmpq $1,%rax
        jne Lwork
    Leval:
        jmp *(%rbx) # encoded in 2-3 bytes.
    <info table>
    Lwork:
        ...

  We could explicitly check for this control flow pattern.

  This is advantageous because:
  * It's optimal if the argument isn't evaluated.
  * If it's evaluated we only have the extra cost of jumping over
    the 2-3 bytes for the call.
  * Guarantees the smaller encoding for the conditional jump.

  However given that Lwork usually has an info table we
  penalize this edge. So Leval should get placed first
  either way and things work out for the best.

  Optimizing for the evaluated case instead would penalize
  the other code path. It adds an jump as we can't fall through
  to Lwork because of the info table.
  Assuming that Lwork is large the chance that the "call" ends up
  in the same cache line is also fairly small.

-}

-- | Look at X number of blocks in two chains to determine
--   if they are "neighbours".
neighbourOverlapp :: Int
neighbourOverlapp = 2

-- | Only edges heavier than this are considered
--   for fusing two chains into a single chain.
fuseEdgeThreshold :: EdgeWeight
fuseEdgeThreshold = 0

-- | Maps blocks near the end of a chain to it's chain AND
-- the other blocks near the end.
-- [A,B,C,D,E] Gives entries like (B -> ([A,B], [A,B,C,D,E]))
-- where [A,B] are blocks in the end region of a chain.
-- This is cheaper then recomputing the ends multiple times.
type FrontierMap = LabelMap ([BlockId],BlockChain)

-- | A non empty ordered sequence of basic blocks.
--   It is suitable for serialization in this order.
--
--   We use OrdList instead of [] to allow fast append on both sides
--   when combining chains.
newtype BlockChain
    = BlockChain { chainBlocks :: (OrdList BlockId) }

instance Eq (BlockChain) where
    (BlockChain blks1) == (BlockChain blks2)
        = fromOL blks1 == fromOL blks2

-- Useful for things like sets and debugging purposes, sorts by blocks
-- in the chain.
instance Ord (BlockChain) where
   (BlockChain lbls1) `compare` (BlockChain lbls2)
       = (fromOL lbls1) `compare` (fromOL lbls2)

instance Outputable (BlockChain) where
    ppr (BlockChain blks) =
        parens (text "Chain:" <+> ppr (fromOL $ blks) )

data WeightedEdge = WeightedEdge !BlockId !BlockId EdgeWeight deriving (Eq)

-- | Non deterministic! (Uniques) Sorts edges by weight and nodes.
instance Ord WeightedEdge where
  compare (WeightedEdge from1 to1 weight1)
          (WeightedEdge from2 to2 weight2)
    | weight1 < weight2 || weight1 == weight2 && from1 < from2 ||
      weight1 == weight2 && from1 == from2 && to1 < to2
    = LT
    | from1 == from2 && to1 == to2 && weight1 == weight2
    = EQ
    | otherwise
    = GT

instance Outputable WeightedEdge where
    ppr (WeightedEdge from to info) =
        ppr from <> text "->" <> ppr to <> brackets (ppr info)

type WeightedEdgeList = [WeightedEdge]

noDups :: [BlockChain] -> Bool
noDups chains =
    let chainBlocks = concatMap chainToBlocks chains :: [BlockId]
        (_blocks, dups) = removeDups compare chainBlocks
    in if null dups then True
        else pprTrace "Duplicates:" (ppr (map toList dups) $$ text "chains" <+> ppr chains ) False

inFront :: BlockId -> BlockChain -> Bool
inFront bid (BlockChain seq)
  = headOL seq == bid

chainMember :: BlockId -> BlockChain -> Bool
chainMember bid chain
  = elem bid $ fromOL . chainBlocks $ chain
--   = setMember bid . chainMembers $ chain

chainSingleton :: BlockId -> BlockChain
chainSingleton lbl
    = BlockChain (unitOL lbl)

chainSnoc :: BlockChain -> BlockId -> BlockChain
chainSnoc (BlockChain blks) lbl
  = BlockChain (blks `snocOL` lbl)

chainConcat :: BlockChain -> BlockChain -> BlockChain
chainConcat (BlockChain blks1) (BlockChain blks2)
  = BlockChain (blks1 `appOL` blks2)

chainToBlocks :: BlockChain -> [BlockId]
chainToBlocks (BlockChain blks) = fromOL blks

-- | Given the Chain A -> B -> C -> D and we break at C
--   we get the two Chains (A -> B, C -> D) as result.
breakChainAt :: BlockId -> BlockChain
             -> (BlockChain,BlockChain)
breakChainAt bid (BlockChain blks)
    | not (bid == head rblks)
    = panic "Block not in chain"
    | otherwise
    = (BlockChain (toOL lblks),
       BlockChain (toOL rblks))
  where
    (lblks, rblks) = break (\lbl -> lbl == bid) (fromOL blks)

takeR :: Int -> BlockChain -> [BlockId]
takeR n (BlockChain blks) =
    take n . fromOLReverse $ blks

takeL :: Int -> BlockChain -> [BlockId]
takeL n (BlockChain blks) =
    take n . fromOL $ blks

-- | For a given list of chains try to fuse chains with strong
--   edges between them into a single chain.
--   Returns the list of fused chains together with a set of
--   used edges. The set of edges is indirectly encoded in the
--   chains so doesn't need to be considered for later passes.
fuseChains :: WeightedEdgeList -> LabelMap BlockChain
           -> (LabelMap BlockChain, Set.Set WeightedEdge)
fuseChains weights chains
    = let fronts = mapFromList $
                    map (\chain -> (headOL . chainBlocks $ chain,chain)) $
                    mapElems chains :: LabelMap BlockChain
          (chains', used, _) = applyEdges weights chains fronts Set.empty
      in (chains', used)
    where
        applyEdges :: WeightedEdgeList -> LabelMap BlockChain
                   -> LabelMap BlockChain -> Set.Set WeightedEdge
                   -> (LabelMap BlockChain, Set.Set WeightedEdge, LabelMap BlockChain)
        applyEdges [] chainsEnd chainsFront used
            = (chainsEnd, used, chainsFront)
        applyEdges (edge@(WeightedEdge from to w):edges) chainsEnd chainsFront used
            --Since we order edges descending by weight we can stop here
            | w <= fuseEdgeThreshold
            = ( chainsEnd, used, chainsFront)
            --Fuse the two chains
            | Just c1 <- mapLookup from chainsEnd
            , Just c2 <- mapLookup to chainsFront
            , c1 /= c2
            = let newChain = chainConcat c1 c2
                  front = headOL . chainBlocks $ newChain
                  end = lastOL . chainBlocks $ newChain
                  chainsFront' = mapInsert front newChain $
                                 mapDelete to chainsFront
                  chainsEnd'   = mapInsert end newChain $
                                 mapDelete from chainsEnd
              in applyEdges edges chainsEnd' chainsFront'
                            (Set.insert edge used)
            | otherwise
            --Check next edge
            = applyEdges edges chainsEnd chainsFront used

-- See also Note [Chain based CFG serialization]
-- We have the chains (A-B-C-D) and (E-F) and an Edge C->E.
--
-- While placing the later after the former doesn't result in sequential
-- control flow it is still be benefical since block C and E might end
-- up in the same cache line.
--
-- So we place these chains next to each other even if we can't fuse them.
--
--   A -> B -> C -> D
--             v
--             - -> E -> F ...
--
-- Simple heuristic to chose which chains we want to combine:
--   * Process edges in descending priority.
--   * Check if there is a edge near the end of one chain which goes
--     to a block near the start of another edge.
--
-- While we could take into account the space between the two blocks which
-- share an edge this blows up compile times quite a bit. It requires
-- us to find all edges between two chains, check the distance for all edges,
-- rank them based on the distance and and only then we can select two chains
-- to combine. Which would add a lot of complexity for little gain.

-- | For a given list of chains and edges try to combine chains with strong
--   edges between them.
combineNeighbourhood :: WeightedEdgeList -> [BlockChain]
                     -> [BlockChain]
combineNeighbourhood edges chains
    = -- pprTraceIt "Neigbours" $
      applyEdges edges endFrontier startFrontier
    where
        --Build maps from chain ends to chains
        endFrontier, startFrontier :: FrontierMap
        endFrontier =
            mapFromList $ concatMap (\chain ->
                                let ends = getEnds chain :: [BlockId]
                                    entry = (ends,chain)
                                in map (\x -> (x,entry)) ends ) chains
        startFrontier =
            mapFromList $ concatMap (\chain ->
                                let front = getFronts chain
                                    entry = (front,chain)
                                in map (\x -> (x,entry)) front) chains
        applyEdges :: WeightedEdgeList -> FrontierMap -> FrontierMap
                   -> [BlockChain]
        applyEdges [] chainEnds _chainFronts =
            ordNub $ map snd $ mapElems chainEnds
        applyEdges ((WeightedEdge from to _w):edges) chainEnds chainFronts
            | Just (c1_e,c1) <- mapLookup from chainEnds
            , Just (c2_f,c2) <- mapLookup to chainFronts
            , c1 /= c2 -- Avoid trying to concat a short chain with itself.
            = let newChain = chainConcat c1 c2
                  newChainFrontier = getFronts newChain
                  newChainEnds = getEnds newChain
                  newFronts :: FrontierMap
                  newFronts =
                    let withoutOld =
                            foldl' (\m b -> mapDelete b m :: FrontierMap) chainFronts (c2_f ++ getFronts c1)
                        entry =
                            (newChainFrontier,newChain) --let bound to ensure sharing
                    in foldl' (\m x -> mapInsert x entry m)
                              withoutOld newChainFrontier

                  newEnds =
                    let withoutOld = foldl' (\m b -> mapDelete b m) chainEnds (c1_e ++ getEnds c2)
                        entry = (newChainEnds,newChain) --let bound to ensure sharing
                    in foldl' (\m x -> mapInsert x entry m)
                              withoutOld newChainEnds
              in
                -- pprTrace "ApplyEdges"
                --  (text "before" $$
                --   text "fronts" <+> ppr chainFronts $$
                --   text "ends" <+> ppr chainEnds $$

                --   text "various" $$
                --   text "newChain" <+> ppr newChain $$
                --   text "newChainFrontier" <+> ppr newChainFrontier $$
                --   text "newChainEnds" <+> ppr newChainEnds $$
                --   text "drop" <+> ppr ((c2_f ++ getFronts c1) ++ (c1_e ++ getEnds c2)) $$

                --   text "after" $$
                --   text "fronts" <+> ppr newFronts $$
                --   text "ends" <+> ppr newEnds
                --   )
                 applyEdges edges newEnds newFronts
            | otherwise
            = --pprTrace "noNeigbours" (ppr ()) $
              applyEdges edges chainEnds chainFronts
         where

        getFronts chain = takeL neighbourOverlapp chain
        getEnds chain = takeR neighbourOverlapp chain

-- See [Chain based CFG serialization]
buildChains :: CFG -> [BlockId]
            -> ( LabelMap BlockChain  -- Resulting chains.
               , Set.Set (BlockId, BlockId)) --List of fused edges.
buildChains succWeights blocks
  = let (_, fusedEdges, chains) = buildNext setEmpty mapEmpty blocks Set.empty
    in (chains, fusedEdges)
  where
    -- We keep a map from the last block in a chain to the chain itself.
    -- This we we can easily check if an block should be appened to an
    -- existing chain!
    buildNext :: LabelSet
              -> LabelMap BlockChain -- Map from last element to chain.
              -> [BlockId] -- Blocks to place
              -> Set.Set (BlockId, BlockId)
              -> ( [BlockChain]  -- Placed Blocks
                 , Set.Set (BlockId, BlockId) --List of fused edges
                 , LabelMap BlockChain
                 )
    buildNext _placed chains [] linked =
        ([], linked, chains)
    buildNext placed chains (block:todo) linked
        | setMember block placed
        = buildNext placed chains todo linked
        | otherwise
        = buildNext placed' chains' todo linked'
      where
        placed' = (foldl' (flip setInsert) placed placedBlocks)
        linked' = Set.union linked linkedEdges
        (placedBlocks, chains', linkedEdges) = findChain block

        --Add the block to a existing or new chain
        --Returns placed blocks, list of resulting chains
        --and fused edges
        findChain :: BlockId
                -> ([BlockId],LabelMap BlockChain, Set.Set (BlockId, BlockId))
        findChain block
        -- B) place block at end of existing chain if
        -- there is no better block to append.
          | (pred:_) <- preds
          , alreadyPlaced pred
          , Just predChain <- mapLookup pred chains
          , (best:_) <- filter (not . alreadyPlaced) $ getSuccs pred
          , best == lbl
          = --pprTrace "B.2)" (ppr (pred,lbl)) $
            let newChain = chainSnoc predChain block
                chainMap = mapInsert lbl newChain $ mapDelete pred chains
            in  ( [lbl]
                , chainMap
                , Set.singleton (pred,lbl) )

          | otherwise
          = --pprTrace "single" (ppr lbl)
            ( [lbl]
            , mapInsert lbl (chainSingleton lbl) chains
            , Set.empty)
            where
              alreadyPlaced blkId = (setMember blkId placed)
              lbl = block
              getSuccs = map fst . getSuccEdgesSorted succWeights
              preds = map fst $ getSuccEdgesSorted predWeights lbl
    --For efficiency we also create the map to look up predecessors here
    predWeights = reverseEdges succWeights

-- We make the CFG a Hoopl Graph, so we can reuse revPostOrder.
newtype BlockNode e x = BN (BlockId,[BlockId])
instance NonLocal (BlockNode) where
  entryLabel (BN (lbl,_))   = lbl
  successors (BN (_,succs)) = succs

fromNode :: BlockNode C C -> BlockId
fromNode (BN x) = fst x

sequenceChain :: forall a i. (Instruction i, Outputable i) => LabelMap a -> CFG
            -> [GenBasicBlock i] -> [GenBasicBlock i]
sequenceChain _info _weights    [] = []
sequenceChain _info _weights    [x] = [x]
sequenceChain  info weights'     blocks@((BasicBlock entry _):_) =
    --Optimization, delete edges of weight <= 0.
    --This significantly improves performance whenever
    --we iterate over all edges, which is a few times!
    let weights :: CFG
        weights
            = filterEdges (\_f _t edgeInfo -> edgeWeight edgeInfo > 0) weights'
        blockMap :: LabelMap (GenBasicBlock i)
        blockMap
            = foldl' (\m blk@(BasicBlock lbl _ins) ->
                        mapInsert lbl blk m)
                     mapEmpty blocks

        toNode :: BlockId -> BlockNode C C
        toNode bid =
            -- sorted such that heavier successors come first.
            BN (bid,map fst . getSuccEdgesSorted weights' $ bid)

        orderedBlocks :: [BlockId]
        orderedBlocks
            = map fromNode $
              revPostorderFrom (fmap (toNode . blockId) blockMap) entry

        (builtChains, builtEdges)
            = {-# SCC "buildChains" #-}
              --pprTraceIt "generatedChains" $
              --pprTrace "orderedBlocks" (ppr orderedBlocks) $
              buildChains weights orderedBlocks

        rankedEdges :: WeightedEdgeList
        -- Sort edges descending, remove fused eges
        rankedEdges =
            map (\(from, to, weight) -> WeightedEdge from to weight) .
            filter (\(from, to, _)
                        -> not (Set.member (from,to) builtEdges)) .
            sortWith (\(_,_,w) -> - w) $ weightedEdgeList weights

        (fusedChains, fusedEdges)
            = ASSERT(noDups $ mapElems builtChains)
              {-# SCC "fuseChains" #-}
              --(pprTrace "RankedEdges" $ ppr rankedEdges) $
              --pprTraceIt "FusedChains" $
              fuseChains rankedEdges builtChains

        rankedEdges' =
            filter (\edge -> not $ Set.member edge fusedEdges) $ rankedEdges

        neighbourChains
            = ASSERT(noDups $ mapElems fusedChains)
              {-# SCC "groupNeighbourChains" #-}
              --pprTraceIt "ResultChains" $
              combineNeighbourhood rankedEdges' (mapElems fusedChains)

        --Make sure the first block stays first
        ([entryChain],chains')
            = ASSERT(noDups $ neighbourChains)
              partition (chainMember entry) neighbourChains
        (entryChain':entryRest)
            | inFront entry entryChain = [entryChain]
            | (rest,entry) <- breakChainAt entry entryChain
            = [entry,rest]
            | otherwise = pprPanic "Entry point eliminated" $
                            ppr ([entryChain],chains')

        prepedChains
            = entryChain':(entryRest++chains') :: [BlockChain]
        blockList
            -- = (concatMap chainToBlocks prepedChains)
            = (concatMap fromOL $ map chainBlocks prepedChains)

        --chainPlaced = setFromList $ map blockId blockList :: LabelSet
        chainPlaced = setFromList $ blockList :: LabelSet
        unplaced =
            let blocks = mapKeys blockMap
                isPlaced b = setMember (b) chainPlaced
            in filter (\block -> not (isPlaced block)) blocks

        placedBlocks =
            --pprTraceIt "placedBlocks" $
            blockList ++ unplaced
        getBlock bid = expectJust "Block placment" $ mapLookup bid blockMap
    in
        --Assert we placed all blocks given as input
        ASSERT(all (\bid -> mapMember bid blockMap) placedBlocks)
        dropJumps info $ map getBlock placedBlocks

dropJumps :: forall a i. Instruction i => LabelMap a -> [GenBasicBlock i]
          -> [GenBasicBlock i]
dropJumps _    [] = []
dropJumps info ((BasicBlock lbl ins):todo)
    | not . null $ ins --This can happen because of shortcutting
    , [dest] <- jumpDestsOfInstr (last ins)
    , ((BasicBlock nextLbl _) : _) <- todo
    , not (mapMember dest info)
    , nextLbl == dest
    = BasicBlock lbl (init ins) : dropJumps info todo
    | otherwise
    = BasicBlock lbl ins : dropJumps info todo

-- -----------------------------------------------------------------------------
-- Sequencing the basic blocks

-- Cmm BasicBlocks are self-contained entities: they always end in a
-- jump, either non-local or to another basic block in the same proc.
-- In this phase, we attempt to place the basic blocks in a sequence
-- such that as many of the local jumps as possible turn into
-- fallthroughs.

sequenceTop
    :: (Instruction instr, Outputable instr)
    => DynFlags --Use new layout code
    -> NcgImpl statics instr jumpDest -> CFG
    -> NatCmmDecl statics instr -> NatCmmDecl statics instr

sequenceTop _     _       _           top@(CmmData _ _) = top
sequenceTop dflags ncgImpl edgeWeights
            (CmmProc info lbl live (ListGraph blocks))
  | (gopt Opt_CfgBlocklayout dflags) && backendMaintainsCfg dflags
  --Use chain based algorithm
  = CmmProc info lbl live ( ListGraph $ ncgMakeFarBranches ncgImpl info $
                            sequenceChain info edgeWeights blocks )
  | otherwise
  --Use old algorithm
  = CmmProc info lbl live ( ListGraph $ ncgMakeFarBranches ncgImpl info $
                            sequenceBlocks cfg info blocks)
  where
    cfg
      | (gopt Opt_WeightlessBlocklayout dflags) ||
        (not $ backendMaintainsCfg dflags)
      -- Don't make use of cfg in the old algorithm
      = Nothing
      -- Use cfg in the old algorithm
      | otherwise = Just edgeWeights

-- The old algorithm:
-- It is very simple (and stupid): We make a graph out of
-- the blocks where there is an edge from one block to another iff the
-- first block ends by jumping to the second.  Then we topologically
-- sort this graph.  Then traverse the list: for each block, we first
-- output the block, then if it has an out edge, we move the
-- destination of the out edge to the front of the list, and continue.

-- FYI, the classic layout for basic blocks uses postorder DFS; this
-- algorithm is implemented in Hoopl.

sequenceBlocks :: Instruction inst => Maybe CFG -> LabelMap a
               -> [GenBasicBlock inst] -> [GenBasicBlock inst]
sequenceBlocks _edgeWeight _ [] = []
sequenceBlocks edgeWeights infos (entry:blocks) =
    let entryNode = mkNode edgeWeights entry
        bodyNodes = reverse
                    (flattenSCCs (sccBlocks edgeWeights blocks))
    in dropJumps infos . seqBlocks infos $ ( entryNode : bodyNodes)
  -- the first block is the entry point ==> it must remain at the start.

sccBlocks
        :: Instruction instr
        => Maybe CFG -> [NatBasicBlock instr]
        -> [SCC (Node BlockId (NatBasicBlock instr))]
sccBlocks edgeWeights blocks =
    stronglyConnCompFromEdgedVerticesUniqR
        (map (mkNode edgeWeights) blocks)

mkNode :: (Instruction t)
       => Maybe CFG -> GenBasicBlock t
       -> Node BlockId (GenBasicBlock t)
mkNode edgeWeights block@(BasicBlock id instrs) =
    DigraphNode block id outEdges
  where
    outEdges :: [BlockId]
    outEdges
      --Select the heaviest successor, ignore weights <= zero
      = successor
      where
        successor
          | Just successors <- fmap (`getSuccEdgesSorted` id)
                                    edgeWeights -- :: Maybe [(Label, EdgeInfo)]
          = case successors of
            [] -> []
            ((target,info):_)
              | length successors > 2 || edgeWeight info <= 0 -> []
              | otherwise -> [target]
          | otherwise
          = case jumpDestsOfInstr (last instrs) of
                [one] -> [one]
                _many -> []

seqBlocks :: LabelMap i -> [Node BlockId (GenBasicBlock t1)]
                        -> [GenBasicBlock t1]
seqBlocks infos blocks = placeNext pullable0 todo0
  where
    -- pullable: Blocks that are not yet placed
    -- todo:     Original order of blocks, to be followed if we have no good
    --           reason not to;
    --           may include blocks that have already been placed, but then
    --           these are not in pullable
    pullable0 = listToUFM [ (i,(b,n)) | DigraphNode b i n <- blocks ]
    todo0     = map node_key blocks

    placeNext _ [] = []
    placeNext pullable (i:rest)
        | Just (block, pullable') <- lookupDeleteUFM pullable i
        = place pullable' rest block
        | otherwise
        -- We already placed this block, so ignore
        = placeNext pullable rest

    place pullable todo (block,[])
                          = block : placeNext pullable todo
    place pullable todo (block@(BasicBlock id instrs),[next])
        | mapMember next infos
        = block : placeNext pullable todo
        | Just (nextBlock, pullable') <- lookupDeleteUFM pullable next
        = BasicBlock id instrs : place pullable' todo nextBlock
        | otherwise
        = block : placeNext pullable todo
    place _ _ (_,tooManyNextNodes)
        = pprPanic "seqBlocks" (ppr tooManyNextNodes)

lookupDeleteUFM :: Uniquable key => UniqFM elt -> key
                -> Maybe (elt, UniqFM elt)
lookupDeleteUFM m k = do -- Maybe monad
    v <- lookupUFM m k
    return (v, delFromUFM m k)


--
-- Copyright (c) 2018 Andreas Klebinger
--

{-# LANGUAGE TypeFamilies, ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE CPP #-}

module CFG
    ( CFG, CfgEdge(..), EdgeInfo(..), EdgeWeight(..)
    , TransitionSource(..)

    --Modify the CFG
    , addWeightEdge, addEdge, delEdge
    , addNodesBetween, shortcutWeightMap
    , reverseEdges, filterEdges
    , addImmediateSuccessor
    , mkWeightInfo, adjustEdgeWeight

    --Query the CFG
    , infoEdgeList, edgeList
    , getSuccessorEdges, getSuccessors
    , getSuccEdgesSorted, weightedEdgeList
    , getEdgeInfo
    , getCfgNodes, hasNode
    , loopMembers

    --Construction/Misc
    , getCfg, getCfgProc, pprEdgeWeights, sanityCheckCfg

    --Find backedges and update their weight
    , optimizeCFG )
where

#include "HsVersions.h"

import GhcPrelude

import BlockId
import Cmm ( RawCmmDecl, GenCmmDecl( .. ), CmmBlock, succ, g_entry
           , CmmGraph )
import CmmNode
import CmmUtils
import CmmSwitch
import Hoopl.Collections
import Hoopl.Label
import Hoopl.Block
import qualified Hoopl.Graph as G

import Util
import Digraph

import Outputable
-- DEBUGGING ONLY
--import Debug
--import OrdList
--import Debug.Trace
import PprCmm ()
import qualified DynFlags as D

import Data.List

-- import qualified Data.IntMap.Strict as M --TODO: LabelMap

type Edge = (BlockId, BlockId)
type Edges = [Edge]

newtype EdgeWeight
  = EdgeWeight Int
  deriving (Eq,Ord,Enum,Num,Real,Integral)

instance Outputable EdgeWeight where
  ppr (EdgeWeight w) = ppr w

type EdgeInfoMap edgeInfo = LabelMap (LabelMap edgeInfo)

-- | A control flow graph where edges have been annotated with a weight.
type CFG = EdgeInfoMap EdgeInfo

data CfgEdge
  = CfgEdge
  { edgeFrom :: !BlockId
  , edgeTo :: !BlockId
  , edgeInfo :: !EdgeInfo
  }

-- | Careful! Since we assume there is at most one edge from A to B
--   the Eq instance does not consider weight.
instance Eq CfgEdge where
  (==) (CfgEdge from1 to1 _) (CfgEdge from2 to2 _)
    = from1 == from2 && to1 == to2

-- | Edges are sorted ascending pointwise by weight, source and destination
instance Ord CfgEdge where
  compare (CfgEdge from1 to1 (EdgeInfo {edgeWeight = weight1}))
          (CfgEdge from2 to2 (EdgeInfo {edgeWeight = weight2}))
    | weight1 < weight2 || weight1 == weight2 && from1 < from2 ||
      weight1 == weight2 && from1 == from2 && to1 < to2
    = LT
    | from1 == from2 && to1 == to2 && weight1 == weight2
    = EQ
    | otherwise
    = GT

instance Outputable CfgEdge where
  ppr (CfgEdge from1 to1 edgeInfo)
    = parens (ppr from1 <+> text "-(" <> ppr edgeInfo <> text ")->" <+> ppr to1)

-- | Can we trace back a edge to a specific Cmm Node
-- or has it been introduced for codegen. We use this to maintain
-- some information which would otherwise be lost during the
-- Cmm <-> asm transition.
-- See also Note [Inverting Conditional Branches]
data TransitionSource
  = CmmSource (CmmNode O C)
  | AsmCodeGen
  deriving (Eq)

-- | Information about edges
data EdgeInfo
  = EdgeInfo
  { transitionSource :: !TransitionSource
  , edgeWeight :: !EdgeWeight
  } deriving (Eq)

instance Outputable EdgeInfo where
  ppr edgeInfo = text "weight:" <+> ppr (edgeWeight edgeInfo)

-- Allow specialization
{-# INLINEABLE mkWeightInfo #-}
-- | Convenience function, generate edge info based
--   on weight not originating from cmm.
mkWeightInfo :: Integral n => n -> EdgeInfo
mkWeightInfo = EdgeInfo AsmCodeGen . fromIntegral

-- | Adjust the weight between the blocks using the given function.
--   If there is no such edge returns the original map.
adjustEdgeWeight :: CFG -> (EdgeWeight -> EdgeWeight)
                 -> BlockId -> BlockId -> CFG
adjustEdgeWeight cfg f from to
  | Just info <- getEdgeInfo from to cfg
  , weight <- edgeWeight info
  = addEdge from to (info { edgeWeight = f weight}) cfg
  | otherwise = cfg

getCfgNodes :: CFG -> LabelSet
getCfgNodes m = mapFoldMapWithKey (\k v -> setFromList (k:mapKeys v)) m

hasNode :: CFG -> BlockId -> Bool
hasNode m node = mapMember node m || any (mapMember node) m

-- | Check if the nodes in the cfg and the set of blocks are the same.
--   In a case of a missmatch we panic and show the difference.
sanityCheckCfg :: CFG -> LabelSet -> SDoc -> Bool
sanityCheckCfg m blockSet msg
    | blockSet == cfgNodes
    = True
    | otherwise =
        pprPanic "Block list and cfg nodes don't match" (
            text "difference:" <+> ppr diff $$
            text "blocks:" <+> ppr blockSet $$
            text "cfg:" <+> ppr m $$
            msg )
            False
    where
      cfgNodes = getCfgNodes m :: LabelSet
      diff = (setUnion cfgNodes blockSet) `setDifference` (setIntersection cfgNodes blockSet) :: LabelSet

-- | Filter the CFG with a custom function f.
--   Paramaeters are `f from to edgeInfo`
filterEdges :: (BlockId -> BlockId -> EdgeInfo -> Bool) -> CFG -> CFG
filterEdges f cfg =
    mapMapWithKey filterSources cfg
    where
      filterSources from m =
        mapFilterWithKey (\to w -> f from to w) m

{- Note [Updating the CFG during shortcutting]

See Note [What is shortcutting] in the control flow optimization
code (CmmContFlowOpt.hs) for a slightly more in depth explanation on shortcutting.

In the native backend we shortcut jumps at the assembly level. (AsmCodeGen.hs)
This means we remove blocks containing only one jump from the code
and instead redirecting all jumps targeting this block to the deleted
blocks jump target.

However we want to have an accurate representation of control
flow in the CFG. So we add/remove edges accordingly to account
for the eliminated blocks and new edges.

If we shortcut A -> B -> C to A -> C:
* We delete edges A -> B and B -> C
* Replacing them with the edge A -> C

We also try to preserve jump weights while doing so.

Note that:
* The edge B -> C can't have interesting weights since
  the block B consists of a single unconditional jump without branching.
* We delete the edge A -> B and add the edge A -> C.
* The edge A -> B can be one of many edges originating from A so likely
  has edge weights we want to preserve.

For this reason we simply store the edge info from the original A -> B
edge and apply this information to the new edge A -> C.

Sometimes we have a scenario where jump target C is not represented by an
BlockId but an immediate value. I'm only aware of this happening without
tables next to code currently.

Then we go from A ---> B - -> IMM   to   A - -> IMM where the dashed arrows
are not stored in the CFG.

In that case we simply delete the edge A -> B.

In terms of implementation the native backend first builds a mapping
from blocks suitable for shortcutting to their jump targets.
Then it redirects all jump instructions to these blocks using the
built up mapping.
This function (shortcutWeightMap) takes the same mapping and
applies the mapping to the CFG in the way layed out above.

-}
shortcutWeightMap :: CFG -> LabelMap (Maybe BlockId) -> CFG
shortcutWeightMap cfg cuts =
  foldl' applyMapping cfg $ mapToList cuts
    where
-- takes the tuple (B,C) from the notation in [Updating the CFG during shortcutting]
      applyMapping :: CFG -> (BlockId,Maybe BlockId) -> CFG
      --Shortcut immediate
      applyMapping m (from, Nothing) =
        mapDelete from .
        fmap (mapDelete from) $ m
      --Regular shortcut
      applyMapping m (from, Just to) =
        let updatedMap :: CFG
            updatedMap
              = fmap (shortcutEdge (from,to)) $
                (mapDelete from m :: CFG )
        --Sometimes we can shortcut multiple blocks like so:
        -- A -> B -> C -> D -> E => A -> E
        -- so we check for such chains.
        in case mapLookup to cuts of
            Nothing -> updatedMap
            Just dest -> applyMapping updatedMap (to, dest)
      --Redirect edge from B to C
      shortcutEdge :: (BlockId, BlockId) -> LabelMap EdgeInfo -> LabelMap EdgeInfo
      shortcutEdge (from, to) m =
        case mapLookup from m of
          Just info -> mapInsert to info $ mapDelete from m
          Nothing   -> m

-- | Sometimes we insert a block which should unconditionally be executed
--   after a given block. This function updates the CFG for these cases.
--  So we get A -> B    => A -> A' -> B
--             \                  \
--              -> C    =>         -> C
--
addImmediateSuccessor :: BlockId -> BlockId -> CFG -> CFG
addImmediateSuccessor node follower cfg
    = updateEdges . addWeightEdge node follower uncondWeight $ cfg
    where
        uncondWeight = fromIntegral . D.uncondWeight .
                       D.cfgWeightInfo $ D.unsafeGlobalDynFlags
        targets = getSuccessorEdges cfg node
        successors = map fst targets :: [BlockId]
        updateEdges = addNewSuccs . remOldSuccs
        remOldSuccs m = foldl' (flip (delEdge node)) m successors
        addNewSuccs m =
          foldl' (\m' (t,info) -> addEdge follower t info m') m targets

-- | Adds a new edge, overwrites existing edges if present
addEdge :: BlockId -> BlockId -> EdgeInfo -> CFG -> CFG
addEdge from to info cfg =
    mapAlter addDest from cfg
    where
        addDest Nothing = Just $ mapSingleton to info
        addDest (Just wm) = Just $ mapInsert to info wm

-- | Adds a edge with the given weight to the cfg
--   If there already existed an edge it is overwritten.
--   `addWeightEdge from to weight cfg`
addWeightEdge :: BlockId -> BlockId -> EdgeWeight -> CFG -> CFG
addWeightEdge from to weight cfg =
    addEdge from to (mkWeightInfo weight) cfg

delEdge :: BlockId -> BlockId -> CFG -> CFG
delEdge from to m =
    mapAlter remDest from m
    where
        remDest Nothing = Nothing
        remDest (Just wm) = Just $ mapDelete to wm

-- | Destinations from bid ordered by weight (descending)
getSuccEdgesSorted :: CFG -> BlockId -> [(BlockId,EdgeInfo)]
getSuccEdgesSorted m bid =
    let destMap = mapFindWithDefault mapEmpty bid m
        cfgEdges = mapToList destMap
        sortedEdges = sortWith (negate . edgeWeight . snd) cfgEdges
    in  --pprTrace "getSuccEdgesSorted" (ppr bid <+> text "map:" <+> ppr m)
        sortedEdges

-- | Get successors of a given node with edge weights.
getSuccessorEdges :: CFG -> BlockId -> [(BlockId,EdgeInfo)]
getSuccessorEdges m bid = maybe [] mapToList $ mapLookup bid m

getEdgeInfo :: BlockId -> BlockId -> CFG -> Maybe EdgeInfo
getEdgeInfo from to m
    | Just wm <- mapLookup from m
    , Just info <- mapLookup to wm
    = Just $! info
    | otherwise
    = Nothing

reverseEdges :: CFG -> CFG
reverseEdges cfg = foldr add mapEmpty flatElems
  where
    elems = mapToList $ fmap mapToList cfg :: [(BlockId,[(BlockId,EdgeInfo)])]
    flatElems =
        concatMap (\(from,ws) -> map (\(to,info) -> (to,from,info)) ws ) elems
    add (to,from,info) m = addEdge to from info m

-- | Returns a unordered list of all edges with info
infoEdgeList :: CFG -> [CfgEdge]
infoEdgeList m =
  mapFoldMapWithKey
    (\from toMap ->
      map (\(to,info) -> CfgEdge from to info) (mapToList toMap))
    m

-- | Unordered list of edges with weight as Tuple (from,to,weight)
weightedEdgeList :: CFG -> [(BlockId,BlockId,EdgeWeight)]
weightedEdgeList m =
  mapFoldMapWithKey
    (\from toMap ->
      map (\(to,info) ->
        (from,to, edgeWeight info)) (mapToList toMap))
    m
      --  (\(from, tos) -> map (\(to,info) -> (from,to, edgeWeight info)) tos )

-- | Returns a unordered list of all edges without weights
edgeList :: CFG -> [Edge]
edgeList m =
        mapFoldMapWithKey (\from toMap -> fmap (from,) (mapKeys toMap)) m

-- | Get successors of a given node without edge weights.
getSuccessors :: CFG -> BlockId -> [BlockId]
getSuccessors m bid
    | Just wm <- mapLookup bid m
    = mapKeys wm
    | otherwise = []

pprEdgeWeights :: CFG -> SDoc
pprEdgeWeights m =
    let edges = sort $ weightedEdgeList m
        printEdge (from, to, weight)
            = text "\t" <> ppr from <+> text "->" <+> ppr to <>
              text "[label=\"" <> ppr weight <> text "\",weight=\"" <>
              ppr weight <> text "\"];\n"
        --for the case that there are no edges from/to this node.
        --This should rarely happen but it can save a lot of time
        --to immediately see it when it does.
        printNode node
            = text "\t" <> ppr node <> text ";\n"
        getEdgeNodes (from, to, _weight) = [from,to]
        edgeNodes = setFromList $ concatMap getEdgeNodes edges :: LabelSet
        nodes = filter (\n -> (not . setMember n) edgeNodes) . mapKeys $ mapFilter null m
    in
    text "digraph {\n" <>
        (foldl' (<>) empty (map printEdge edges)) <>
        (foldl' (<>) empty (map printNode nodes)) <>
    text "}\n"

{-# INLINE updateEdgeWeight #-} --Allows eliminating the tuple when possible
updateEdgeWeight :: (EdgeWeight -> EdgeWeight) -> Edge -> CFG -> CFG
updateEdgeWeight f (from, to) cfg
    | Just oldInfo <- getEdgeInfo from to cfg
    = let oldWeight = edgeWeight oldInfo
          newWeight = f oldWeight
      in addEdge from to (oldInfo {edgeWeight = newWeight}) cfg
    | otherwise
    = panic "Trying to update invalid edge"

-- from to oldWeight => newWeight
mapWeights :: (BlockId -> BlockId -> EdgeWeight -> EdgeWeight) -> CFG -> CFG
mapWeights f cfg =
  foldl' (\cfg (CfgEdge from to info) ->
            let oldWeight = edgeWeight info
                newWeight = f from to oldWeight
            in addEdge from to (info {edgeWeight = newWeight}) cfg)
          cfg (infoEdgeList cfg)

-- | Insert a block in the control flow between two other blocks.
-- We pass a list of tuples (A,B,C) where
-- * A -> C: Old edge
-- * A -> B -> C : New Arc, where B is the new block.
-- It's possible that a block has two jumps to the same block
-- in the assembly code. However we still only store a single edge for
-- these cases.
-- We assign the old edge info to the edge A -> B and assign B -> C the
-- weight of an unconditional jump.
addNodesBetween :: CFG -> [(BlockId,BlockId,BlockId)] -> CFG
addNodesBetween m updates =
  foldl'  updateWeight m .
          weightUpdates $ updates
    where
      weight = fromIntegral . D.uncondWeight .
                D.cfgWeightInfo $ D.unsafeGlobalDynFlags
      -- We might add two blocks for different jumps along a single
      -- edge. So we end up with edges:   A -> B -> C   ,   A -> D -> C
      -- in this case after applying the first update the weight for A -> C
      -- is no longer available. So we calculate future weights before updates.
      weightUpdates = map getWeight
      getWeight :: (BlockId,BlockId,BlockId) -> (BlockId,BlockId,BlockId,EdgeInfo)
      getWeight (from,between,old)
        | Just edgeInfo <- getEdgeInfo from old m
        = (from,between,old,edgeInfo)
        | otherwise
        = pprPanic "Can't find weight for edge that should have one" (
            text "triple" <+> ppr (from,between,old) $$
            text "updates" <+> ppr updates )
      updateWeight :: CFG -> (BlockId,BlockId,BlockId,EdgeInfo) -> CFG
      updateWeight m (from,between,old,edgeInfo)
        = addEdge from between edgeInfo .
          addWeightEdge between old weight .
          delEdge from old $ m

{-
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
  ~~~       Note [CFG Edge Weights]    ~~~
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

  Edge weights assigned do not currently represent a specific
  cost model and rather just a ranking of which blocks should
  be placed next to each other given their connection type in
  the CFG.
  This is especially relevant if we whenever two blocks will
  jump to the same target.

                     A   B
                      \ /
                       C

  Should A or B be placed in front of C? The block layout algorithm
  decides this based on which edge (A,C)/(B,C) is heavier. So we
  make a educated guess how often execution will transer control
  along each edge as well as how much we gain by placing eg A before
  C.

  We rank edges in this order:
  * Unconditional Control Transfer - They will always
    transfer control to their target. Unless there is a info table
    we can turn the jump into a fallthrough as well.
    We use 20k as default, so it's easy to spot if values have been
    modified but unlikely that we run into issues with overflow.
  * If branches (likely) - We assume branches marked as likely
    are taken more than 80% of the time.
    By ranking them below unconditional jumps we make sure we
    prefer the unconditional if there is a conditional and
    unconditional edge towards a block.
  * If branches (regular) - The false branch can potentially be turned
    into a fallthrough so we prefer it slightly over the true branch.
  * Unlikely branches - These can be assumed to be taken less than 20%
    of the time. So we given them one of the lowest priorities.
  * Switches - Switches at this level are implemented as jump tables
    so have a larger number of successors. So without more information
    we can only say that each individual successor is unlikely to be
    jumped to and we rank them accordingly.
  * Calls - We currently ignore calls completly:
        * By the time we return from a call there is a good chance
          that the address we return to has already been evicted from
          cache eliminating a main advantage sequential placement brings.
        * Calls always require a info table in front of their return
          address. This reduces the chance that we return to the same
          cache line further.

-}
-- | Generate weights for a Cmm proc based on some simple heuristics.
getCfgProc :: D.CfgWeights -> RawCmmDecl -> CFG
getCfgProc _       (CmmData {}) = mapEmpty
getCfgProc weights (CmmProc _info _lab _live graph) = getCfg weights graph

getCfg :: D.CfgWeights -> CmmGraph -> CFG
getCfg weights graph =
  foldl' insertEdge edgelessCfg $ concatMap getBlockEdges blocks
  where
    D.CFGWeights
            { D.uncondWeight = uncondWeight
            , D.condBranchWeight = condBranchWeight
            , D.switchWeight = switchWeight
            , D.callWeight = callWeight
            , D.likelyCondWeight = likelyCondWeight
            , D.unlikelyCondWeight = unlikelyCondWeight
            --  Last two are used in other places
            --, D.infoTablePenalty = infoTablePenalty
            --, D.backEdgeBonus = backEdgeBonus
            } = weights
    -- Explicitly add all nodes to the cfg to ensure they are part of the
    -- CFG.
    edgelessCfg = mapFromList $ zip (map G.entryLabel blocks) (repeat mapEmpty)
    insertEdge :: CFG -> ((BlockId,BlockId),EdgeInfo) -> CFG
    insertEdge m ((from,to),weight) =
      mapAlter f from m
        where
          f :: Maybe (LabelMap EdgeInfo) -> Maybe (LabelMap EdgeInfo)
          f Nothing = Just $ mapSingleton to weight
          f (Just destMap) = Just $ mapInsert to weight destMap
    getBlockEdges :: CmmBlock -> [((BlockId,BlockId),EdgeInfo)]
    getBlockEdges block =
      case branch of
        CmmBranch dest -> [mkEdge dest uncondWeight]
        CmmCondBranch _c t f l
          | l == Nothing ->
              [mkEdge f condBranchWeight,   mkEdge t condBranchWeight]
          | l == Just True ->
              [mkEdge f unlikelyCondWeight, mkEdge t likelyCondWeight]
          | l == Just False ->
              [mkEdge f likelyCondWeight,   mkEdge t unlikelyCondWeight]
        (CmmSwitch _e ids) ->
          let switchTargets = switchTargetsToList ids
              --Compiler performance hack - for very wide switches don't
              --consider targets for layout.
              adjustedWeight =
                if (length switchTargets > 10) then -1 else switchWeight
          in map (\x -> mkEdge x adjustedWeight) switchTargets
        (CmmCall { cml_cont = Just cont})  -> [mkEdge cont callWeight]
        (CmmForeignCall {Cmm.succ = cont}) -> [mkEdge cont callWeight]
        (CmmCall { cml_cont = Nothing })   -> []
        other ->
            panic "Foo" $
            ASSERT2(False, ppr "Unkown successor cause:" <>
              (ppr branch <+> text "=>" <> ppr (G.successors other)))
            map (\x -> ((bid,x),mkEdgeInfo 0)) $ G.successors other
      where
        bid = G.entryLabel block
        mkEdgeInfo = EdgeInfo (CmmSource branch) . fromIntegral
        mkEdge target weight = ((bid,target), mkEdgeInfo weight)
        branch = lastNode block :: CmmNode O C

    blocks = revPostorder graph :: [CmmBlock]

--Find back edges by BFS
findBackEdges :: BlockId -> CFG -> Edges
findBackEdges root cfg =
    --pprTraceIt "Backedges:" $
    map fst .
    filter (\x -> snd x == Backward) $ typedEdges
  where
    edges = edgeList cfg :: [(BlockId,BlockId)]
    getSuccs = getSuccessors cfg :: BlockId -> [BlockId]
    typedEdges =
      classifyEdges root getSuccs edges :: [((BlockId,BlockId),EdgeType)]

optimizeCFG :: D.CfgWeights -> RawCmmDecl -> CFG -> CFG
optimizeCFG _ (CmmData {}) cfg = cfg
optimizeCFG weights (CmmProc info _lab _live graph) cfg =
    favourFewerPreds  .
    penalizeInfoTables info .
    increaseBackEdgeWeight (g_entry graph) $ cfg
  where

    -- | Increase the weight of all backedges in the CFG
    -- this helps to make loop jumpbacks the heaviest edges
    increaseBackEdgeWeight :: BlockId -> CFG -> CFG
    increaseBackEdgeWeight root cfg =
        let backedges = findBackEdges root cfg
            update weight
              --Keep irrelevant edges irrelevant
              | weight <= 0 = 0
              | otherwise
              = weight + fromIntegral (D.backEdgeBonus weights)
        in  foldl'  (\cfg edge -> updateEdgeWeight update edge cfg)
                    cfg backedges

    -- | Since we cant fall through info tables we penalize these.
    penalizeInfoTables :: LabelMap a -> CFG -> CFG
    penalizeInfoTables info cfg =
        mapWeights fupdate cfg
      where
        fupdate :: BlockId -> BlockId -> EdgeWeight -> EdgeWeight
        fupdate _ to weight
          | mapMember to info
          = weight - (fromIntegral $ D.infoTablePenalty weights)
          | otherwise = weight

{- Note [Optimize for Fallthrough]

-}
    -- | If a block has two successors, favour the one with fewer
    -- predecessors. (As that one is more likely to become a fallthrough)
    favourFewerPreds :: CFG -> CFG
    favourFewerPreds cfg =
        let
            revCfg =
              reverseEdges $ filterEdges
                              (\_from -> fallthroughTarget)  cfg

            predCount n = length $ getSuccessorEdges revCfg n
            nodes = getCfgNodes cfg

            modifiers :: Int -> Int -> (EdgeWeight, EdgeWeight)
            modifiers preds1 preds2
              | preds1 <  preds2 = ( 1,-1)
              | preds1 == preds2 = ( 0, 0)
              | otherwise        = (-1, 1)

            update cfg node
              | [(s1,e1),(s2,e2)] <- getSuccessorEdges cfg node
              , w1 <- edgeWeight e1
              , w2 <- edgeWeight e2
              --Only change the weights if there isn't already a ordering.
              , w1 == w2
              , (mod1,mod2) <- modifiers (predCount s1) (predCount s2)
              = (\cfg' ->
                  (adjustEdgeWeight cfg' (+mod2) node s2))
                  (adjustEdgeWeight cfg  (+mod1) node s1)
              | otherwise
              = cfg
        in setFoldl update cfg nodes
      where
        fallthroughTarget :: BlockId -> EdgeInfo -> Bool
        fallthroughTarget to (EdgeInfo source _weight)
          | mapMember to info = False
          | AsmCodeGen <- source = True
          | CmmSource (CmmBranch {}) <- source = True
          | CmmSource (CmmCondBranch {}) <- source = True
          | otherwise = False

-- | Determine loop membership of blocks based on SCC analysis
--   Ideally we would replace this with a variant giving us loop
--   levels instead but the SCC code will do for now.
loopMembers :: CFG -> LabelMap Bool
loopMembers cfg =
    foldl' (flip setLevel) mapEmpty sccs
  where
    mkNode :: BlockId -> Node BlockId BlockId
    mkNode bid = DigraphNode bid bid (getSuccessors cfg bid)
    nodes = map mkNode (setElems $ getCfgNodes cfg)

    sccs = stronglyConnCompFromEdgedVerticesOrd nodes

    setLevel :: SCC BlockId -> LabelMap Bool -> LabelMap Bool
    setLevel (AcyclicSCC bid) m = mapInsert bid False m
    setLevel (CyclicSCC bids) m = foldl' (\m k -> mapInsert k True m) m bids


-- | Generating C symbol names emitted by the compiler.
module CPrim
    ( atomicReadLabel
    , atomicWriteLabel
    , atomicRMWLabel
    , cmpxchgLabel
    , popCntLabel
    , pdepLabel
    , pextLabel
    , bSwapLabel
    , clzLabel
    , ctzLabel
    , word2FloatLabel
    ) where

import GhcPrelude

import CmmType
import CmmMachOp
import Outputable

popCntLabel :: Width -> String
popCntLabel w = "hs_popcnt" ++ pprWidth w
  where
    pprWidth W8  = "8"
    pprWidth W16 = "16"
    pprWidth W32 = "32"
    pprWidth W64 = "64"
    pprWidth w   = pprPanic "popCntLabel: Unsupported word width " (ppr w)

pdepLabel :: Width -> String
pdepLabel w = "hs_pdep" ++ pprWidth w
  where
    pprWidth W8  = "8"
    pprWidth W16 = "16"
    pprWidth W32 = "32"
    pprWidth W64 = "64"
    pprWidth w   = pprPanic "pdepLabel: Unsupported word width " (ppr w)

pextLabel :: Width -> String
pextLabel w = "hs_pext" ++ pprWidth w
  where
    pprWidth W8  = "8"
    pprWidth W16 = "16"
    pprWidth W32 = "32"
    pprWidth W64 = "64"
    pprWidth w   = pprPanic "pextLabel: Unsupported word width " (ppr w)

bSwapLabel :: Width -> String
bSwapLabel w = "hs_bswap" ++ pprWidth w
  where
    pprWidth W16 = "16"
    pprWidth W32 = "32"
    pprWidth W64 = "64"
    pprWidth w   = pprPanic "bSwapLabel: Unsupported word width " (ppr w)

clzLabel :: Width -> String
clzLabel w = "hs_clz" ++ pprWidth w
  where
    pprWidth W8  = "8"
    pprWidth W16 = "16"
    pprWidth W32 = "32"
    pprWidth W64 = "64"
    pprWidth w   = pprPanic "clzLabel: Unsupported word width " (ppr w)

ctzLabel :: Width -> String
ctzLabel w = "hs_ctz" ++ pprWidth w
  where
    pprWidth W8  = "8"
    pprWidth W16 = "16"
    pprWidth W32 = "32"
    pprWidth W64 = "64"
    pprWidth w   = pprPanic "ctzLabel: Unsupported word width " (ppr w)

word2FloatLabel :: Width -> String
word2FloatLabel w = "hs_word2float" ++ pprWidth w
  where
    pprWidth W32 = "32"
    pprWidth W64 = "64"
    pprWidth w   = pprPanic "word2FloatLabel: Unsupported word width " (ppr w)

atomicRMWLabel :: Width -> AtomicMachOp -> String
atomicRMWLabel w amop = "hs_atomic_" ++ pprFunName amop ++ pprWidth w
  where
    pprWidth W8  = "8"
    pprWidth W16 = "16"
    pprWidth W32 = "32"
    pprWidth W64 = "64"
    pprWidth w   = pprPanic "atomicRMWLabel: Unsupported word width " (ppr w)

    pprFunName AMO_Add  = "add"
    pprFunName AMO_Sub  = "sub"
    pprFunName AMO_And  = "and"
    pprFunName AMO_Nand = "nand"
    pprFunName AMO_Or   = "or"
    pprFunName AMO_Xor  = "xor"

cmpxchgLabel :: Width -> String
cmpxchgLabel w = "hs_cmpxchg" ++ pprWidth w
  where
    pprWidth W8  = "8"
    pprWidth W16 = "16"
    pprWidth W32 = "32"
    pprWidth W64 = "64"
    pprWidth w   = pprPanic "cmpxchgLabel: Unsupported word width " (ppr w)

atomicReadLabel :: Width -> String
atomicReadLabel w = "hs_atomicread" ++ pprWidth w
  where
    pprWidth W8  = "8"
    pprWidth W16 = "16"
    pprWidth W32 = "32"
    pprWidth W64 = "64"
    pprWidth w   = pprPanic "atomicReadLabel: Unsupported word width " (ppr w)

atomicWriteLabel :: Width -> String
atomicWriteLabel w = "hs_atomicwrite" ++ pprWidth w
  where
    pprWidth W8  = "8"
    pprWidth W16 = "16"
    pprWidth W32 = "32"
    pprWidth W64 = "64"
    pprWidth w   = pprPanic "atomicWriteLabel: Unsupported word width " (ppr w)


module Dwarf (
  dwarfGen
  ) where

import GhcPrelude

import CLabel
import CmmExpr         ( GlobalReg(..) )
import Config          ( cProjectName, cProjectVersion )
import CoreSyn         ( Tickish(..) )
import Debug
import DynFlags
import Module
import Outputable
import Platform
import Unique
import UniqSupply

import Dwarf.Constants
import Dwarf.Types

import Control.Arrow    ( first )
import Control.Monad    ( mfilter )
import Data.Maybe
import Data.List        ( sortBy )
import Data.Ord         ( comparing )
import qualified Data.Map as Map
import System.FilePath
import System.Directory ( getCurrentDirectory )

import qualified Hoopl.Label as H
import qualified Hoopl.Collections as H

-- | Generate DWARF/debug information
dwarfGen :: DynFlags -> ModLocation -> UniqSupply -> [DebugBlock]
            -> IO (SDoc, UniqSupply)
dwarfGen _  _      us [] = return (empty, us)
dwarfGen df modLoc us blocks = do

  -- Convert debug data structures to DWARF info records
  -- We strip out block information when running with -g0 or -g1.
  let procs = debugSplitProcs blocks
      stripBlocks dbg
        | debugLevel df < 2 = dbg { dblBlocks = [] }
        | otherwise         = dbg
  compPath <- getCurrentDirectory
  let lowLabel = dblCLabel $ head procs
      highLabel = mkAsmTempEndLabel $ dblCLabel $ last procs
      dwarfUnit = DwarfCompileUnit
        { dwChildren = map (procToDwarf df) (map stripBlocks procs)
        , dwName = fromMaybe "" (ml_hs_file modLoc)
        , dwCompDir = addTrailingPathSeparator compPath
        , dwProducer = cProjectName ++ " " ++ cProjectVersion
        , dwLowLabel = lowLabel
        , dwHighLabel = highLabel
        , dwLineLabel = dwarfLineLabel
        }

  -- Check whether we have any source code information, so we do not
  -- end up writing a pointer to an empty .debug_line section
  -- (dsymutil on Mac Os gets confused by this).
  let haveSrcIn blk = isJust (dblSourceTick blk) && isJust (dblPosition blk)
                      || any haveSrcIn (dblBlocks blk)
      haveSrc = any haveSrcIn procs

  -- .debug_abbrev section: Declare the format we're using
  let abbrevSct = pprAbbrevDecls haveSrc

  -- .debug_info section: Information records on procedures and blocks
  let -- unique to identify start and end compilation unit .debug_inf
      (unitU, us') = takeUniqFromSupply us
      infoSct = vcat [ ptext dwarfInfoLabel <> colon
                     , dwarfInfoSection
                     , compileUnitHeader unitU
                     , pprDwarfInfo haveSrc dwarfUnit
                     , compileUnitFooter unitU
                     ]

  -- .debug_line section: Generated mainly by the assembler, but we
  -- need to label it
  let lineSct = dwarfLineSection $$
                ptext dwarfLineLabel <> colon

  -- .debug_frame section: Information about the layout of the GHC stack
  let (framesU, us'') = takeUniqFromSupply us'
      frameSct = dwarfFrameSection $$
                 ptext dwarfFrameLabel <> colon $$
                 pprDwarfFrame (debugFrame framesU procs)

  -- .aranges section: Information about the bounds of compilation units
  let aranges' | gopt Opt_SplitSections df = map mkDwarfARange procs
               | otherwise                 = [DwarfARange lowLabel highLabel]
  let aranges = dwarfARangesSection $$ pprDwarfARanges aranges' unitU

  return (infoSct $$ abbrevSct $$ lineSct $$ frameSct $$ aranges, us'')

-- | Build an address range entry for one proc.
-- With split sections, each proc needs its own entry, since they may get
-- scattered in the final binary. Without split sections, we could make a
-- single arange based on the first/last proc.
mkDwarfARange :: DebugBlock -> DwarfARange
mkDwarfARange proc = DwarfARange start end
  where
    start = dblCLabel proc
    end = mkAsmTempEndLabel start

-- | Header for a compilation unit, establishing global format
-- parameters
compileUnitHeader :: Unique -> SDoc
compileUnitHeader unitU = sdocWithPlatform $ \plat ->
  let cuLabel = mkAsmTempLabel unitU  -- sits right before initialLength field
      length = ppr (mkAsmTempEndLabel cuLabel) <> char '-' <> ppr cuLabel
               <> text "-4"       -- length of initialLength field
  in vcat [ ppr cuLabel <> colon
          , text "\t.long " <> length  -- compilation unit size
          , pprHalf 3                          -- DWARF version
          , sectionOffset (ptext dwarfAbbrevLabel) (ptext dwarfAbbrevLabel)
                                               -- abbrevs offset
          , text "\t.byte " <> ppr (platformWordSize plat) -- word size
          ]

-- | Compilation unit footer, mainly establishing size of debug sections
compileUnitFooter :: Unique -> SDoc
compileUnitFooter unitU =
  let cuEndLabel = mkAsmTempEndLabel $ mkAsmTempLabel unitU
  in ppr cuEndLabel <> colon

-- | Splits the blocks by procedures. In the result all nested blocks
-- will come from the same procedure as the top-level block. See
-- Note [Splitting DebugBlocks] for details.
debugSplitProcs :: [DebugBlock] -> [DebugBlock]
debugSplitProcs b = concat $ H.mapElems $ mergeMaps $ map (split Nothing) b
  where mergeMaps = foldr (H.mapUnionWithKey (const (++))) H.mapEmpty
        split :: Maybe DebugBlock -> DebugBlock -> H.LabelMap [DebugBlock]
        split parent blk = H.mapInsert prc [blk'] nested
          where prc = dblProcedure blk
                blk' = blk { dblBlocks = own_blks
                           , dblParent = parent
                           }
                own_blks = fromMaybe [] $ H.mapLookup prc nested
                nested = mergeMaps $ map (split parent') $ dblBlocks blk
                -- Figure out who should be the parent of nested blocks.
                -- If @blk@ is optimized out then it isn't a good choice
                -- and we just use its parent.
                parent'
                  | Nothing <- dblPosition blk = parent
                  | otherwise                  = Just blk

{-
Note [Splitting DebugBlocks]

DWARF requires that we break up the nested DebugBlocks produced from
the C-- AST. For instance, we begin with tick trees containing nested procs.
For example,

    proc A [tick1, tick2]
      block B [tick3]
        proc C [tick4]

when producing DWARF we need to procs (which are represented in DWARF as
TAG_subprogram DIEs) to be top-level DIEs. debugSplitProcs is responsible for
this transform, pulling out the nested procs into top-level procs.

However, in doing this we need to be careful to preserve the parentage of the
nested procs. This is the reason DebugBlocks carry the dblParent field, allowing
us to reorganize the above tree as,

    proc A [tick1, tick2]
      block B [tick3]
    proc C [tick4] parent=B

Here we have annotated the new proc C with an attribute giving its original
parent, B.
-}

-- | Generate DWARF info for a procedure debug block
procToDwarf :: DynFlags -> DebugBlock -> DwarfInfo
procToDwarf df prc
  = DwarfSubprogram { dwChildren = map (blockToDwarf df) (dblBlocks prc)
                    , dwName     = case dblSourceTick prc of
                         Just s@SourceNote{} -> sourceName s
                         _otherwise -> showSDocDump df $ ppr $ dblLabel prc
                    , dwLabel    = dblCLabel prc
                    , dwParent   = fmap mkAsmTempDieLabel
                                   $ mfilter goodParent
                                   $ fmap dblCLabel (dblParent prc)
                    }
  where
  goodParent a | a == dblCLabel prc = False
               -- Omit parent if it would be self-referential
  goodParent a | not (externallyVisibleCLabel a)
               , debugLevel df < 2 = False
               -- We strip block information when running -g0 or -g1, don't
               -- refer to blocks in that case. Fixes #14894.
  goodParent _ = True

-- | Generate DWARF info for a block
blockToDwarf :: DynFlags -> DebugBlock -> DwarfInfo
blockToDwarf df blk
  = DwarfBlock { dwChildren = concatMap (tickToDwarf df) (dblTicks blk)
                              ++ map (blockToDwarf df) (dblBlocks blk)
               , dwLabel    = dblCLabel blk
               , dwMarker   = marker
               }
  where
    marker
      | Just _ <- dblPosition blk = Just $ mkAsmTempLabel $ dblLabel blk
      | otherwise                 = Nothing   -- block was optimized out

tickToDwarf :: DynFlags -> Tickish () -> [DwarfInfo]
tickToDwarf _  (SourceNote ss _) = [DwarfSrcNote ss]
tickToDwarf _ _ = []

-- | Generates the data for the debug frame section, which encodes the
-- desired stack unwind behaviour for the debugger
debugFrame :: Unique -> [DebugBlock] -> DwarfFrame
debugFrame u procs
  = DwarfFrame { dwCieLabel = mkAsmTempLabel u
               , dwCieInit  = initUws
               , dwCieProcs = map (procToFrame initUws) procs
               }
  where
    initUws :: UnwindTable
    initUws = Map.fromList [(Sp, Just (UwReg Sp 0))]

-- | Generates unwind information for a procedure debug block
procToFrame :: UnwindTable -> DebugBlock -> DwarfFrameProc
procToFrame initUws blk
  = DwarfFrameProc { dwFdeProc    = dblCLabel blk
                   , dwFdeHasInfo = dblHasInfoTbl blk
                   , dwFdeBlocks  = map (uncurry blockToFrame)
                                        (setHasInfo blockUws)
                   }
  where blockUws :: [(DebugBlock, [UnwindPoint])]
        blockUws = map snd $ sortBy (comparing fst) $ flatten blk

        flatten :: DebugBlock
                -> [(Int, (DebugBlock, [UnwindPoint]))]
        flatten b@DebugBlock{ dblPosition=pos, dblUnwind=uws, dblBlocks=blocks }
          | Just p <- pos  = (p, (b, uws')):nested
          | otherwise      = nested -- block was optimized out
          where uws'   = addDefaultUnwindings initUws uws
                nested = concatMap flatten blocks

        -- | If the current procedure has an info table, then we also say that
        -- its first block has one to ensure that it gets the necessary -1
        -- offset applied to its start address.
        -- See Note [Info Offset] in Dwarf.Types.
        setHasInfo :: [(DebugBlock, [UnwindPoint])]
                   -> [(DebugBlock, [UnwindPoint])]
        setHasInfo [] = []
        setHasInfo (c0:cs) = first setIt c0 : cs
          where
            setIt child =
              child { dblHasInfoTbl = dblHasInfoTbl child
                                      || dblHasInfoTbl blk }

blockToFrame :: DebugBlock -> [UnwindPoint] -> DwarfFrameBlock
blockToFrame blk uws
  = DwarfFrameBlock { dwFdeBlkHasInfo = dblHasInfoTbl blk
                    , dwFdeUnwind     = uws
                    }

addDefaultUnwindings :: UnwindTable -> [UnwindPoint] -> [UnwindPoint]
addDefaultUnwindings tbl pts =
    [ UnwindPoint lbl (tbl' `mappend` tbl)
      -- mappend is left-biased
    | UnwindPoint lbl tbl' <- pts
    ]


-- | Formats on this architecture
--      A Format is a combination of width and class
--
--      TODO:   Signed vs unsigned?
--
--      TODO:   This module is currenly shared by all architectures because
--              NCGMonad need to know about it to make a VReg. It would be better
--              to have architecture specific formats, and do the overloading
--              properly. eg SPARC doesn't care about FF80.
--
module Format (
    Format(..),
    intFormat,
    floatFormat,
    isFloatFormat,
    cmmTypeFormat,
    formatToWidth,
    formatInBytes
)

where

import GhcPrelude

import Cmm
import Outputable

-- It looks very like the old MachRep, but it's now of purely local
-- significance, here in the native code generator.  You can change it
-- without global consequences.
--
-- A major use is as an opcode qualifier; thus the opcode
--      mov.l a b
-- might be encoded
--      MOV II32 a b
-- where the Format field encodes the ".l" part.

-- ToDo: it's not clear to me that we need separate signed-vs-unsigned formats
--        here.  I've removed them from the x86 version, we'll see what happens --SDM

-- ToDo: quite a few occurrences of Format could usefully be replaced by Width

data Format
        = II8
        | II16
        | II32
        | II64
        | FF32
        | FF64
        | FF80
        deriving (Show, Eq)

-- | Get the integer format of this width.
intFormat :: Width -> Format
intFormat width
 = case width of
        W8      -> II8
        W16     -> II16
        W32     -> II32
        W64     -> II64
        other   -> sorry $ "The native code generator cannot " ++
            "produce code for Format.intFormat " ++ show other
            ++ "\n\tConsider using the llvm backend with -fllvm"

-- | Get the float format of this width.
floatFormat :: Width -> Format
floatFormat width
 = case width of
        W32     -> FF32
        W64     -> FF64
        W80     -> FF80
        other   -> pprPanic "Format.floatFormat" (ppr other)

-- | Check if a format represents a floating point value.
isFloatFormat :: Format -> Bool
isFloatFormat format
 = case format of
        FF32    -> True
        FF64    -> True
        FF80    -> True
        _       -> False

-- | Convert a Cmm type to a Format.
cmmTypeFormat :: CmmType -> Format
cmmTypeFormat ty
        | isFloatType ty        = floatFormat (typeWidth ty)
        | otherwise             = intFormat (typeWidth ty)

-- | Get the Width of a Format.
formatToWidth :: Format -> Width
formatToWidth format
 = case format of
        II8             -> W8
        II16            -> W16
        II32            -> W32
        II64            -> W64
        FF32            -> W32
        FF64            -> W64
        FF80            -> W80

formatInBytes :: Format -> Int
formatInBytes = widthInBytes . formatToWidth


module Instruction (
        RegUsage(..),
        noUsage,
        GenBasicBlock(..), blockId,
        ListGraph(..),
        NatCmm,
        NatCmmDecl,
        NatBasicBlock,
        topInfoTable,
        entryBlocks,
        Instruction(..)
)

where

import GhcPrelude

import Reg

import BlockId
import Hoopl.Collections
import Hoopl.Label
import DynFlags
import Cmm hiding (topInfoTable)
import Platform

-- | Holds a list of source and destination registers used by a
--      particular instruction.
--
--   Machine registers that are pre-allocated to stgRegs are filtered
--      out, because they are uninteresting from a register allocation
--      standpoint.  (We wouldn't want them to end up on the free list!)
--
--   As far as we are concerned, the fixed registers simply don't exist
--      (for allocation purposes, anyway).
--
data RegUsage
        = RU [Reg] [Reg]

-- | No regs read or written to.
noUsage :: RegUsage
noUsage  = RU [] []

-- Our flavours of the Cmm types
-- Type synonyms for Cmm populated with native code
type NatCmm instr
        = GenCmmGroup
                CmmStatics
                (LabelMap CmmStatics)
                (ListGraph instr)

type NatCmmDecl statics instr
        = GenCmmDecl
                statics
                (LabelMap CmmStatics)
                (ListGraph instr)

type NatBasicBlock instr
        = GenBasicBlock instr

-- | Returns the info table associated with the CmmDecl's entry point,
-- if any.
topInfoTable :: GenCmmDecl a (LabelMap i) (ListGraph b) -> Maybe i
topInfoTable (CmmProc infos _ _ (ListGraph (b:_)))
  = mapLookup (blockId b) infos
topInfoTable _
  = Nothing

-- | Return the list of BlockIds in a CmmDecl that are entry points
-- for this proc (i.e. they may be jumped to from outside this proc).
entryBlocks :: GenCmmDecl a (LabelMap i) (ListGraph b) -> [BlockId]
entryBlocks (CmmProc info _ _ (ListGraph code)) = entries
  where
        infos = mapKeys info
        entries = case code of
                    [] -> infos
                    BasicBlock entry _ : _ -- first block is the entry point
                       | entry `elem` infos -> infos
                       | otherwise          -> entry : infos
entryBlocks _ = []

-- | Common things that we can do with instructions, on all architectures.
--      These are used by the shared parts of the native code generator,
--      specifically the register allocators.
--
class   Instruction instr where

        -- | Get the registers that are being used by this instruction.
        --      regUsage doesn't need to do any trickery for jumps and such.
        --      Just state precisely the regs read and written by that insn.
        --      The consequences of control flow transfers, as far as register
        --      allocation goes, are taken care of by the register allocator.
        --
        regUsageOfInstr
                :: Platform
                -> instr
                -> RegUsage

        -- | Apply a given mapping to all the register references in this
        --      instruction.
        patchRegsOfInstr
                :: instr
                -> (Reg -> Reg)
                -> instr

        -- | Checks whether this instruction is a jump/branch instruction.
        --      One that can change the flow of control in a way that the
        --      register allocator needs to worry about.
        isJumpishInstr
                :: instr -> Bool

        -- | Give the possible destinations of this jump instruction.
        --      Must be defined for all jumpish instructions.
        jumpDestsOfInstr
                :: instr -> [BlockId]

        -- | Change the destination of this jump instruction.
        --      Used in the linear allocator when adding fixup blocks for join
        --      points.
        patchJumpInstr
                :: instr
                -> (BlockId -> BlockId)
                -> instr

        -- | An instruction to spill a register into a spill slot.
        mkSpillInstr
                :: DynFlags
                -> Reg          -- ^ the reg to spill
                -> Int          -- ^ the current stack delta
                -> Int          -- ^ spill slot to use
                -> instr

        -- | An instruction to reload a register from a spill slot.
        mkLoadInstr
                :: DynFlags
                -> Reg          -- ^ the reg to reload.
                -> Int          -- ^ the current stack delta
                -> Int          -- ^ the spill slot to use
                -> instr

        -- | See if this instruction is telling us the current C stack delta
        takeDeltaInstr
                :: instr
                -> Maybe Int

        -- | Check whether this instruction is some meta thing inserted into
        --      the instruction stream for other purposes.
        --
        --      Not something that has to be treated as a real machine instruction
        --      and have its registers allocated.
        --
        --      eg, comments, delta, ldata, etc.
        isMetaInstr
                :: instr
                -> Bool

        -- | Copy the value in a register to another one.
        --      Must work for all register classes.
        mkRegRegMoveInstr
                :: Platform
                -> Reg          -- ^ source register
                -> Reg          -- ^ destination register
                -> instr

        -- | Take the source and destination from this reg -> reg move instruction
        --      or Nothing if it's not one
        takeRegRegMoveInstr
                :: instr
                -> Maybe (Reg, Reg)

        -- | Make an unconditional jump instruction.
        --      For architectures with branch delay slots, its ok to put
        --      a NOP after the jump. Don't fill the delay slot with an
        --      instruction that references regs or you'll confuse the
        --      linear allocator.
        mkJumpInstr
                :: BlockId
                -> [instr]

        -- Subtract an amount from the C stack pointer
        mkStackAllocInstr
                :: Platform
                -> Int
                -> [instr]

        -- Add an amount to the C stack pointer
        mkStackDeallocInstr
                :: Platform
                -> Int
                -> [instr]


{-# LANGUAGE CPP #-}

-- -----------------------------------------------------------------------------
--
-- (c) The University of Glasgow 1993-2004
--
-- The native code generator's monad.
--
-- -----------------------------------------------------------------------------

module NCGMonad (
        NcgImpl(..),
        NatM_State(..), mkNatM_State,

        NatM, -- instance Monad
        initNat,
        addImportNat,
        addNodeBetweenNat,
        addImmediateSuccessorNat,
        updateCfgNat,
        getUniqueNat,
        mapAccumLNat,
        setDeltaNat,
        getDeltaNat,
        getThisModuleNat,
        getBlockIdNat,
        getNewLabelNat,
        getNewRegNat,
        getNewRegPairNat,
        getPicBaseMaybeNat,
        getPicBaseNat,
        getDynFlags,
        getModLoc,
        getFileId,
        getDebugBlock,

        DwarfFiles
)

where

#include "HsVersions.h"

import GhcPrelude

import Reg
import Format
import TargetReg

import BlockId
import Hoopl.Collections
import Hoopl.Label
import CLabel           ( CLabel )
import Debug
import FastString       ( FastString )
import UniqFM
import UniqSupply
import Unique           ( Unique )
import DynFlags
import Module

import Control.Monad    ( liftM, ap )

import Instruction
import Outputable (SDoc, pprPanic, ppr)
import Cmm (RawCmmDecl, CmmStatics)
import CFG

data NcgImpl statics instr jumpDest = NcgImpl {
    cmmTopCodeGen             :: RawCmmDecl -> NatM [NatCmmDecl statics instr],
    generateJumpTableForInstr :: instr -> Maybe (NatCmmDecl statics instr),
    getJumpDestBlockId        :: jumpDest -> Maybe BlockId,
    canShortcut               :: instr -> Maybe jumpDest,
    shortcutStatics           :: (BlockId -> Maybe jumpDest) -> statics -> statics,
    shortcutJump              :: (BlockId -> Maybe jumpDest) -> instr -> instr,
    pprNatCmmDecl             :: NatCmmDecl statics instr -> SDoc,
    maxSpillSlots             :: Int,
    allocatableRegs           :: [RealReg],
    ncg_x86fp_kludge          :: [NatCmmDecl statics instr] -> [NatCmmDecl statics instr],
    ncgExpandTop              :: [NatCmmDecl statics instr] -> [NatCmmDecl statics instr],
    ncgAllocMoreStack         :: Int -> NatCmmDecl statics instr
                              -> UniqSM (NatCmmDecl statics instr, [(BlockId,BlockId)]),
    -- ^ The list of block ids records the redirected jumps to allow us to update
    -- the CFG.
    ncgMakeFarBranches        :: LabelMap CmmStatics -> [NatBasicBlock instr] -> [NatBasicBlock instr],
    extractUnwindPoints       :: [instr] -> [UnwindPoint],
    -- ^ given the instruction sequence of a block, produce a list of
    -- the block's 'UnwindPoint's
    -- See Note [What is this unwinding business?] in Debug
    -- and Note [Unwinding information in the NCG] in this module.
    invertCondBranches        :: CFG -> LabelMap CmmStatics -> [NatBasicBlock instr]
                              -> [NatBasicBlock instr]
    -- ^ Turn the sequence of `jcc l1; jmp l2` into `jncc l2; <block_l1>`
    -- when possible.
    }

data NatM_State
        = NatM_State {
                natm_us          :: UniqSupply,
                natm_delta       :: Int,
                natm_imports     :: [(CLabel)],
                natm_pic         :: Maybe Reg,
                natm_dflags      :: DynFlags,
                natm_this_module :: Module,
                natm_modloc      :: ModLocation,
                natm_fileid      :: DwarfFiles,
                natm_debug_map   :: LabelMap DebugBlock,
                natm_cfg         :: CFG
        -- ^ Having a CFG with additional information is essential for some
        -- operations. However we can't reconstruct all information once we
        -- generated instructions. So instead we update the CFG as we go.
        }

type DwarfFiles = UniqFM (FastString, Int)

newtype NatM result = NatM (NatM_State -> (result, NatM_State))

unNat :: NatM a -> NatM_State -> (a, NatM_State)
unNat (NatM a) = a

mkNatM_State :: UniqSupply -> Int -> DynFlags -> Module -> ModLocation ->
                DwarfFiles -> LabelMap DebugBlock -> CFG -> NatM_State
mkNatM_State us delta dflags this_mod
        = \loc dwf dbg cfg ->
                NatM_State
                        { natm_us = us
                        , natm_delta = delta
                        , natm_imports = []
                        , natm_pic = Nothing
                        , natm_dflags = dflags
                        , natm_this_module = this_mod
                        , natm_modloc = loc
                        , natm_fileid = dwf
                        , natm_debug_map = dbg
                        , natm_cfg = cfg
                        }

initNat :: NatM_State -> NatM a -> (a, NatM_State)
initNat init_st m
        = case unNat m init_st of { (r,st) -> (r,st) }

instance Functor NatM where
      fmap = liftM

instance Applicative NatM where
      pure = returnNat
      (<*>) = ap

instance Monad NatM where
  (>>=) = thenNat

instance MonadUnique NatM where
  getUniqueSupplyM = NatM $ \st ->
      case splitUniqSupply (natm_us st) of
          (us1, us2) -> (us1, st {natm_us = us2})

  getUniqueM = NatM $ \st ->
      case takeUniqFromSupply (natm_us st) of
          (uniq, us') -> (uniq, st {natm_us = us'})

thenNat :: NatM a -> (a -> NatM b) -> NatM b
thenNat expr cont
        = NatM $ \st -> case unNat expr st of
                        (result, st') -> unNat (cont result) st'

returnNat :: a -> NatM a
returnNat result
        = NatM $ \st ->  (result, st)

mapAccumLNat :: (acc -> x -> NatM (acc, y))
                -> acc
                -> [x]
                -> NatM (acc, [y])

mapAccumLNat _ b []
  = return (b, [])
mapAccumLNat f b (x:xs)
  = do (b__2, x__2)  <- f b x
       (b__3, xs__2) <- mapAccumLNat f b__2 xs
       return (b__3, x__2:xs__2)

getUniqueNat :: NatM Unique
getUniqueNat = NatM $ \ st ->
    case takeUniqFromSupply $ natm_us st of
    (uniq, us') -> (uniq, st {natm_us = us'})

instance HasDynFlags NatM where
    getDynFlags = NatM $ \ st -> (natm_dflags st, st)

getDeltaNat :: NatM Int
getDeltaNat = NatM $ \ st -> (natm_delta st, st)

setDeltaNat :: Int -> NatM ()
setDeltaNat delta = NatM $ \ st -> ((), st {natm_delta = delta})

getThisModuleNat :: NatM Module
getThisModuleNat = NatM $ \ st -> (natm_this_module st, st)

addImportNat :: CLabel -> NatM ()
addImportNat imp
        = NatM $ \ st -> ((), st {natm_imports = imp : natm_imports st})

updateCfgNat :: (CFG -> CFG) -> NatM ()
updateCfgNat f
        = NatM $ \ st -> ((), st { natm_cfg = f (natm_cfg st) })

-- | Record that we added a block between `from` and `old`.
addNodeBetweenNat :: BlockId -> BlockId -> BlockId -> NatM ()
addNodeBetweenNat from between to
 = do   df <- getDynFlags
        let jmpWeight = fromIntegral . uncondWeight .
                        cfgWeightInfo $ df
        updateCfgNat (updateCfg jmpWeight from between to)
  where
    -- When transforming A -> B to A -> A' -> B
    -- A -> A' keeps the old edge info while
    -- A' -> B gets the info for an unconditional
    -- jump.
    updateCfg weight from between old m
        | Just info <- getEdgeInfo from old m
        = addEdge from between info .
          addWeightEdge between old weight .
          delEdge from old $ m
        | otherwise
        = pprPanic "Faild to update cfg: Untracked edge" (ppr (from,to))

-- | Place `succ` after `block` and change any edges
--   block -> X to `succ` -> X
addImmediateSuccessorNat :: BlockId -> BlockId -> NatM ()
addImmediateSuccessorNat block succ
        = updateCfgNat (addImmediateSuccessor block succ)

getBlockIdNat :: NatM BlockId
getBlockIdNat
 = do   u <- getUniqueNat
        return (mkBlockId u)

getNewLabelNat :: NatM CLabel
getNewLabelNat
 = blockLbl <$> getBlockIdNat

getNewRegNat :: Format -> NatM Reg
getNewRegNat rep
 = do u <- getUniqueNat
      dflags <- getDynFlags
      return (RegVirtual $ targetMkVirtualReg (targetPlatform dflags) u rep)

getNewRegPairNat :: Format -> NatM (Reg,Reg)
getNewRegPairNat rep
 = do u <- getUniqueNat
      dflags <- getDynFlags
      let vLo = targetMkVirtualReg (targetPlatform dflags) u rep
      let lo  = RegVirtual $ targetMkVirtualReg (targetPlatform dflags) u rep
      let hi  = RegVirtual $ getHiVirtualRegFromLo vLo
      return (lo, hi)

getPicBaseMaybeNat :: NatM (Maybe Reg)
getPicBaseMaybeNat
        = NatM (\state -> (natm_pic state, state))

getPicBaseNat :: Format -> NatM Reg
getPicBaseNat rep
 = do   mbPicBase <- getPicBaseMaybeNat
        case mbPicBase of
                Just picBase -> return picBase
                Nothing
                 -> do
                        reg <- getNewRegNat rep
                        NatM (\state -> (reg, state { natm_pic = Just reg }))

getModLoc :: NatM ModLocation
getModLoc
        = NatM $ \ st -> (natm_modloc st, st)

getFileId :: FastString -> NatM Int
getFileId f = NatM $ \st ->
  case lookupUFM (natm_fileid st) f of
    Just (_,n) -> (n, st)
    Nothing    -> let n = 1 + sizeUFM (natm_fileid st)
                      fids = addToUFM (natm_fileid st) f (f,n)
                  in n `seq` fids `seq` (n, st { natm_fileid = fids  })

getDebugBlock :: Label -> NatM (Maybe DebugBlock)
getDebugBlock l = NatM $ \st -> (mapLookup l (natm_debug_map st), st)