module CmmLive (
        CmmLive,
        BlockEntryLiveness,
        cmmLiveness,
        cmmFormalsToLiveLocals,
  ) where

#include "HsVersions.h"

import BlockId
import Cmm
import Dataflow

import Maybes
import Panic
import UniqSet

-----------------------------------------------------------------------------
-- Calculating what variables are live on entry to a basic block
-----------------------------------------------------------------------------

-- | The variables live on entry to a block
type CmmLive = UniqSet LocalReg

-- | A mapping from block labels to the variables live on entry
type BlockEntryLiveness = BlockEnv CmmLive

-- | A mapping from block labels to the blocks that target it
type BlockSources = BlockEnv (UniqSet BlockId)

-- | A mapping from block labels to the statements in the block
type BlockStmts = BlockEnv [CmmStmt]

-----------------------------------------------------------------------------
-- | Calculated liveness info for a list of 'CmmBasicBlock'
-----------------------------------------------------------------------------
cmmLiveness :: [CmmBasicBlock] -> BlockEntryLiveness
cmmLiveness blocks =
    fixedpoint (cmmBlockDependants sources)
               (cmmBlockUpdate blocks')
               (map blockId blocks)
               (mkBlockEnv [(blockId b, emptyUniqSet) | b <- blocks])
    where
      sources :: BlockSources
      sources = cmmBlockSources blocks

      blocks' :: BlockStmts
      blocks' = mkBlockEnv $ map block_name blocks

      block_name :: CmmBasicBlock -> (BlockId, [CmmStmt])
      block_name b = (blockId b, blockStmts b)

{-
-- For debugging, annotate each block with a comment indicating
-- the calculated live variables
cmmLivenessComment ::
    BlockEnv (UniqSet LocalReg) -> CmmBasicBlock -> CmmBasicBlock
cmmLivenessComment live (BasicBlock ident stmts) =
    BasicBlock ident stmts' where
        stmts' = (CmmComment $ mkFastString $ showSDoc $ ppr $ live'):stmts
        live' = map CmmLocal $ uniqSetToList $ lookupWithDefaultUFM live emptyUniqSet ident
-}


-----------------------------------------------------------------------------
-- | Calculates a table of where one can lookup the blocks that might
-- need updating after a given block is updated in the liveness analysis
-----------------------------------------------------------------------------
cmmBlockSources :: [CmmBasicBlock] -> BlockSources
cmmBlockSources blocks = foldr aux emptyBlockEnv blocks
    where
      aux :: CmmBasicBlock
          -> BlockSources
          -> BlockSources
      aux block sourcesUFM =
          foldUniqSet (add_source_edges $ blockId block)
                      sourcesUFM
                      (branch_targets $ blockStmts block)

      add_source_edges :: BlockId -> BlockId
                       -> BlockSources
                       -> BlockSources
      add_source_edges source target ufm =
          addToBEnv_Acc (flip addOneToUniqSet) unitUniqSet ufm target source

      branch_targets :: [CmmStmt] -> UniqSet BlockId
      branch_targets stmts =
          mkUniqSet $ concatMap target stmts where
              target (CmmBranch ident) = [ident]
              target (CmmCondBranch _ ident) = [ident]
              target (CmmSwitch _ blocks) = mapMaybe id blocks
              target _ = []

-----------------------------------------------------------------------------
-- | Given the table calculated by 'cmmBlockSources', list all blocks
-- that depend on the result of a particular block.
--
-- Used by the call to 'fixedpoint'.
-----------------------------------------------------------------------------
cmmBlockDependants :: BlockSources -> BlockId -> [BlockId]
cmmBlockDependants sources ident =
    uniqSetToList $ lookupWithDefaultBEnv sources emptyUniqSet ident

-----------------------------------------------------------------------------
-- | Given the table of type 'BlockStmts' and a block that was updated,
-- calculate an updated BlockEntryLiveness
-----------------------------------------------------------------------------
cmmBlockUpdate ::
    BlockStmts
    -> BlockId
    -> Maybe BlockId
    -> BlockEntryLiveness
    -> Maybe BlockEntryLiveness
cmmBlockUpdate blocks node _ state =
    if (sizeUniqSet old_live) == (sizeUniqSet new_live)
      then Nothing
      else Just $ extendBlockEnv state node new_live
    where
      new_live, old_live :: CmmLive
      new_live = cmmStmtListLive state block_stmts
      old_live = lookupWithDefaultBEnv state missing_live node

      block_stmts :: [CmmStmt]
      block_stmts = lookupWithDefaultBEnv blocks missing_block node

      missing_live = panic "unknown block id during liveness analysis"
      missing_block = panic "unknown block id during liveness analysis"

-----------------------------------------------------------------------------
-- Section: 
-----------------------------------------------------------------------------
-----------------------------------------------------------------------------
-- CmmBlockLive, cmmStmtListLive and helpers
-----------------------------------------------------------------------------

-- Calculate the live registers for a local block (list of statements)

cmmStmtListLive :: BlockEntryLiveness -> [CmmStmt] -> CmmLive
cmmStmtListLive other_live stmts =
    foldr ((.) . (cmmStmtLive other_live)) id stmts emptyUniqSet

-----------------------------------------------------------------------------
-- This code is written in the style of a state monad,
-- but since Control.Monad.State is not in the core
-- we can't use it in GHC, so we'll fake one here.
-- We don't need a return value so well leave it out.
-- Thus 'bind' reduces to function composition.

type CmmLivenessTransformer = CmmLive -> CmmLive

-- Helpers for the "Monad"
addLive, addKilled :: CmmLive -> CmmLivenessTransformer
addLive new_live live = live `unionUniqSets` new_live
addKilled new_killed live = live `minusUniqSet` new_killed

--------------------------------
-- Liveness of a CmmStmt
--------------------------------
cmmFormalsToLiveLocals :: HintedCmmFormals -> [LocalReg]
cmmFormalsToLiveLocals formals = map hintlessCmm formals

cmmStmtLive :: BlockEntryLiveness -> CmmStmt -> CmmLivenessTransformer
cmmStmtLive _ (CmmNop) = id
cmmStmtLive _ (CmmComment _) = id
cmmStmtLive _ (CmmAssign reg expr) =
    cmmExprLive expr . reg_liveness where
        reg_liveness =
            case reg of
              (CmmLocal reg') -> addKilled $ unitUniqSet reg'
              (CmmGlobal _) -> id
cmmStmtLive _ (CmmStore expr1 expr2) =
    cmmExprLive expr2 . cmmExprLive expr1
cmmStmtLive _ (CmmCall target results arguments _ _) =
    target_liveness .
    foldr ((.) . cmmExprLive) id (map hintlessCmm arguments) .
    addKilled (mkUniqSet $ cmmFormalsToLiveLocals results) where
        target_liveness =
            case target of
              (CmmCallee target _) -> cmmExprLive target
              (CmmPrim _) -> id
cmmStmtLive other_live (CmmBranch target) =
    addLive (lookupWithDefaultBEnv other_live emptyUniqSet target)
cmmStmtLive other_live (CmmCondBranch expr target) =
    cmmExprLive expr .
    addLive (lookupWithDefaultBEnv other_live emptyUniqSet target)
cmmStmtLive other_live (CmmSwitch expr targets) =
    cmmExprLive expr .
    (foldr ((.) . (addLive .
                   lookupWithDefaultBEnv other_live emptyUniqSet))
           id
           (mapCatMaybes id targets))
cmmStmtLive _ (CmmJump expr params) =
    const (cmmExprLive expr $ foldr ((.) . cmmExprLive) id (map hintlessCmm params) $ emptyUniqSet)
cmmStmtLive _ (CmmReturn params) =
    const (foldr ((.) . cmmExprLive) id (map hintlessCmm params) $ emptyUniqSet)

--------------------------------
-- Liveness of a CmmExpr
--------------------------------
cmmExprLive :: CmmExpr -> CmmLivenessTransformer
cmmExprLive expr = addLive (mkUniqSet $ expr_liveness expr) where
    expr_liveness :: CmmExpr -> [LocalReg]
    expr_liveness (CmmLit _) = []
    expr_liveness (CmmLoad expr _) = expr_liveness expr
    expr_liveness (CmmReg reg) = reg_liveness reg
    expr_liveness (CmmMachOp _ exprs) = concatMap expr_liveness exprs
    expr_liveness (CmmRegOff reg _) = reg_liveness reg
    expr_liveness (CmmStackSlot _ _) = panic "cmmExprLive CmmStackSlot"

    reg_liveness :: CmmReg -> [LocalReg]
    reg_liveness (CmmLocal reg) = [reg]
    reg_liveness (CmmGlobal _) = []