module StgCmmMonad (
        FCode,        
        initC, runC, fixC,
        newUnique,
        emitLabel,
        emit, emitDecl, emitProc,
        emitProcWithConvention, emitProcWithStackFrame,
        emitOutOfLine, emitAssign, emitStore,
        emitComment, emitTick, emitUnwind,
        getCmm, aGraphToGraph,
        getCodeR, getCode, getCodeScoped, getHeapUsage,
        mkCmmIfThenElse, mkCmmIfThen, mkCmmIfGoto,
        mkCmmIfThenElse', mkCmmIfThen', mkCmmIfGoto',
        mkCall, mkCmmCall,
        forkClosureBody, forkLneBody, forkAlts, codeOnly,
        ConTagZ,
        Sequel(..), ReturnKind(..),
        withSequel, getSequel,
        setTickyCtrLabel, getTickyCtrLabel,
        tickScope, getTickScope,
        withUpdFrameOff, getUpdFrameOff, initUpdFrameOff,
        HeapUsage(..), VirtualHpOffset,        initHpUsage,
        getHpUsage,  setHpUsage, heapHWM,
        setVirtHp, getVirtHp, setRealHp,
        getModuleName,
        
        getState, setState, getSelfLoop, withSelfLoop, getInfoDown, getDynFlags, getThisPackage,
        
        CgIdInfo(..),
        getBinds, setBinds,
        
        CgInfoDownwards(..), CgState(..)        
    ) where
#include "HsVersions.h"
import GhcPrelude hiding( sequence, succ )
import Cmm
import StgCmmClosure
import DynFlags
import Hoopl.Collections
import MkGraph
import BlockId
import CLabel
import SMRep
import Module
import Id
import VarEnv
import OrdList
import BasicTypes( ConTagZ )
import Unique
import UniqSupply
import FastString
import Outputable
import Control.Monad
import Data.List
newtype FCode a = FCode { doFCode :: CgInfoDownwards -> CgState -> (a, CgState) }
instance Functor FCode where
    fmap f (FCode g) = FCode $ \i s -> case g i s of (a, s') -> (f a, s')
instance Applicative FCode where
    pure val = FCode (\_info_down state -> (val, state))
    
    (<*>) = ap
instance Monad FCode where
    FCode m >>= k = FCode $
        \info_down state ->
            case m info_down state of
              (m_result, new_state) ->
                 case k m_result of
                   FCode kcode -> kcode info_down new_state
    
instance MonadUnique FCode where
  getUniqueSupplyM = cgs_uniqs <$> getState
  getUniqueM = FCode $ \_ st ->
    let (u, us') = takeUniqFromSupply (cgs_uniqs st)
    in (u, st { cgs_uniqs = us' })
initC :: IO CgState
initC  = do { uniqs <- mkSplitUniqSupply 'c'
            ; return (initCgState uniqs) }
runC :: DynFlags -> Module -> CgState -> FCode a -> (a,CgState)
runC dflags mod st fcode = doFCode fcode (initCgInfoDown dflags mod) st
fixC :: (a -> FCode a) -> FCode a
fixC fcode = FCode $
    \info_down state -> let (v, s) = doFCode (fcode v) info_down state
                        in (v, s)
data CgInfoDownwards        
  = MkCgInfoDown {
        cgd_dflags    :: DynFlags,
        cgd_mod       :: Module,            
        cgd_updfr_off :: UpdFrameOffset,    
        cgd_ticky     :: CLabel,            
        cgd_sequel    :: Sequel,            
        cgd_self_loop :: Maybe SelfLoopInfo,
                                            
                                            
                                            
        cgd_tick_scope:: CmmTickScope       
  }
type CgBindings = IdEnv CgIdInfo
data CgIdInfo
  = CgIdInfo
        { cg_id :: Id   
                        
                        
                        
        , cg_lf  :: LambdaFormInfo
        , cg_loc :: CgLoc                     
        }
instance Outputable CgIdInfo where
  ppr (CgIdInfo { cg_id = id, cg_loc = loc })
    = ppr id <+> text "-->" <+> ppr loc
data Sequel
  = Return              
  | AssignTo
        [LocalReg]      
                        
                        
        Bool            
                        
                        
                        
                        
instance Outputable Sequel where
    ppr Return = text "Return"
    ppr (AssignTo regs b) = text "AssignTo" <+> ppr regs <+> ppr b
data ReturnKind
  = AssignedDirectly
  | ReturnedTo BlockId ByteOff
initCgInfoDown :: DynFlags -> Module -> CgInfoDownwards
initCgInfoDown dflags mod
  = MkCgInfoDown { cgd_dflags    = dflags
                 , cgd_mod       = mod
                 , cgd_updfr_off = initUpdFrameOff dflags
                 , cgd_ticky     = mkTopTickyCtrLabel
                 , cgd_sequel    = initSequel
                 , cgd_self_loop = Nothing
                 , cgd_tick_scope= GlobalScope }
initSequel :: Sequel
initSequel = Return
initUpdFrameOff :: DynFlags -> UpdFrameOffset
initUpdFrameOff dflags = widthInBytes (wordWidth dflags) 
data CgState
  = MkCgState {
     cgs_stmts :: CmmAGraph,          
     cgs_tops  :: OrdList CmmDecl,
        
        
        
     cgs_binds :: CgBindings,
     cgs_hp_usg  :: HeapUsage,
     cgs_uniqs :: UniqSupply }
data HeapUsage   
  = HeapUsage {
        virtHp :: VirtualHpOffset,       
                                         
        realHp :: VirtualHpOffset        
                                         
    }
type VirtualHpOffset = WordOff
initCgState :: UniqSupply -> CgState
initCgState uniqs
  = MkCgState { cgs_stmts  = mkNop
              , cgs_tops   = nilOL
              , cgs_binds  = emptyVarEnv
              , cgs_hp_usg = initHpUsage
              , cgs_uniqs  = uniqs }
stateIncUsage :: CgState -> CgState -> CgState
stateIncUsage s1 s2@(MkCgState { cgs_hp_usg = hp_usg })
     = s1 { cgs_hp_usg  = cgs_hp_usg  s1 `maxHpHw`  virtHp hp_usg }
       `addCodeBlocksFrom` s2
addCodeBlocksFrom :: CgState -> CgState -> CgState
s1 `addCodeBlocksFrom` s2
  = s1 { cgs_stmts = cgs_stmts s1 MkGraph.<*> cgs_stmts s2,
         cgs_tops  = cgs_tops  s1 `appOL` cgs_tops  s2 }
heapHWM :: HeapUsage -> VirtualHpOffset
heapHWM = virtHp
initHpUsage :: HeapUsage
initHpUsage = HeapUsage { virtHp = 0, realHp = 0 }
maxHpHw :: HeapUsage -> VirtualHpOffset -> HeapUsage
hp_usg `maxHpHw` hw = hp_usg { virtHp = virtHp hp_usg `max` hw }
getState :: FCode CgState
getState = FCode $ \_info_down state -> (state, state)
setState :: CgState -> FCode ()
setState state = FCode $ \_info_down _ -> ((), state)
getHpUsage :: FCode HeapUsage
getHpUsage = do
        state <- getState
        return $ cgs_hp_usg state
setHpUsage :: HeapUsage -> FCode ()
setHpUsage new_hp_usg = do
        state <- getState
        setState $ state {cgs_hp_usg = new_hp_usg}
setVirtHp :: VirtualHpOffset -> FCode ()
setVirtHp new_virtHp
  = do  { hp_usage <- getHpUsage
        ; setHpUsage (hp_usage {virtHp = new_virtHp}) }
getVirtHp :: FCode VirtualHpOffset
getVirtHp
  = do  { hp_usage <- getHpUsage
        ; return (virtHp hp_usage) }
setRealHp ::  VirtualHpOffset -> FCode ()
setRealHp new_realHp
  = do  { hp_usage <- getHpUsage
        ; setHpUsage (hp_usage {realHp = new_realHp}) }
getBinds :: FCode CgBindings
getBinds = do
        state <- getState
        return $ cgs_binds state
setBinds :: CgBindings -> FCode ()
setBinds new_binds = do
        state <- getState
        setState $ state {cgs_binds = new_binds}
withState :: FCode a -> CgState -> FCode (a,CgState)
withState (FCode fcode) newstate = FCode $ \info_down state ->
  case fcode info_down newstate of
    (retval, state2) -> ((retval,state2), state)
newUniqSupply :: FCode UniqSupply
newUniqSupply = do
        state <- getState
        let (us1, us2) = splitUniqSupply (cgs_uniqs state)
        setState $ state { cgs_uniqs = us1 }
        return us2
newUnique :: FCode Unique
newUnique = do
        state <- getState
        let (u,us') = takeUniqFromSupply (cgs_uniqs state)
        setState $ state { cgs_uniqs = us' }
        return u
getInfoDown :: FCode CgInfoDownwards
getInfoDown = FCode $ \info_down state -> (info_down,state)
getSelfLoop :: FCode (Maybe SelfLoopInfo)
getSelfLoop = do
        info_down <- getInfoDown
        return $ cgd_self_loop info_down
withSelfLoop :: SelfLoopInfo -> FCode a -> FCode a
withSelfLoop self_loop code = do
        info_down <- getInfoDown
        withInfoDown code (info_down {cgd_self_loop = Just self_loop})
instance HasDynFlags FCode where
    getDynFlags = liftM cgd_dflags getInfoDown
getThisPackage :: FCode UnitId
getThisPackage = liftM thisPackage getDynFlags
withInfoDown :: FCode a -> CgInfoDownwards -> FCode a
withInfoDown (FCode fcode) info_down = FCode $ \_ state -> fcode info_down state
getModuleName :: FCode Module
getModuleName = do { info <- getInfoDown; return (cgd_mod info) }
withSequel :: Sequel -> FCode a -> FCode a
withSequel sequel code
  = do  { info  <- getInfoDown
        ; withInfoDown code (info {cgd_sequel = sequel, cgd_self_loop = Nothing }) }
getSequel :: FCode Sequel
getSequel = do  { info <- getInfoDown
                ; return (cgd_sequel info) }
withUpdFrameOff :: UpdFrameOffset -> FCode a -> FCode a
withUpdFrameOff size code
  = do  { info  <- getInfoDown
        ; withInfoDown code (info {cgd_updfr_off = size }) }
getUpdFrameOff :: FCode UpdFrameOffset
getUpdFrameOff
  = do  { info  <- getInfoDown
        ; return $ cgd_updfr_off info }
getTickyCtrLabel :: FCode CLabel
getTickyCtrLabel = do
        info <- getInfoDown
        return (cgd_ticky info)
setTickyCtrLabel :: CLabel -> FCode a -> FCode a
setTickyCtrLabel ticky code = do
        info <- getInfoDown
        withInfoDown code (info {cgd_ticky = ticky})
getTickScope :: FCode CmmTickScope
getTickScope = do
        info <- getInfoDown
        return (cgd_tick_scope info)
tickScope :: FCode a -> FCode a
tickScope code = do
        info <- getInfoDown
        if debugLevel (cgd_dflags info) == 0 then code else do
          u <- newUnique
          let scope' = SubScope u (cgd_tick_scope info)
          withInfoDown code info{ cgd_tick_scope = scope' }
forkClosureBody :: FCode () -> FCode ()
forkClosureBody body_code
  = do  { dflags <- getDynFlags
        ; info   <- getInfoDown
        ; us     <- newUniqSupply
        ; state  <- getState
        ; let body_info_down = info { cgd_sequel    = initSequel
                                    , cgd_updfr_off = initUpdFrameOff dflags
                                    , cgd_self_loop = Nothing }
              fork_state_in = (initCgState us) { cgs_binds = cgs_binds state }
              ((),fork_state_out) = doFCode body_code body_info_down fork_state_in
        ; setState $ state `addCodeBlocksFrom` fork_state_out }
forkLneBody :: FCode a -> FCode a
forkLneBody body_code
  = do  { info_down <- getInfoDown
        ; us        <- newUniqSupply
        ; state     <- getState
        ; let fork_state_in = (initCgState us) { cgs_binds = cgs_binds state }
              (result, fork_state_out) = doFCode body_code info_down fork_state_in
        ; setState $ state `addCodeBlocksFrom` fork_state_out
        ; return result }
codeOnly :: FCode () -> FCode ()
codeOnly body_code
  = do  { info_down <- getInfoDown
        ; us        <- newUniqSupply
        ; state     <- getState
        ; let   fork_state_in = (initCgState us) { cgs_binds   = cgs_binds state
                                                 , cgs_hp_usg  = cgs_hp_usg state }
                ((), fork_state_out) = doFCode body_code info_down fork_state_in
        ; setState $ state `addCodeBlocksFrom` fork_state_out }
forkAlts :: [FCode a] -> FCode [a]
forkAlts branch_fcodes
  = do  { info_down <- getInfoDown
        ; us <- newUniqSupply
        ; state <- getState
        ; let compile us branch
                = (us2, doFCode branch info_down branch_state)
                where
                  (us1,us2) = splitUniqSupply us
                  branch_state = (initCgState us1) {
                                        cgs_binds  = cgs_binds state
                                      , cgs_hp_usg = cgs_hp_usg state }
              (_us, results) = mapAccumL compile us branch_fcodes
              (branch_results, branch_out_states) = unzip results
        ; setState $ foldl stateIncUsage state branch_out_states
                
        ; return branch_results }
getCodeR :: FCode a -> FCode (a, CmmAGraph)
getCodeR fcode
  = do  { state1 <- getState
        ; (a, state2) <- withState fcode (state1 { cgs_stmts = mkNop })
        ; setState $ state2 { cgs_stmts = cgs_stmts state1  }
        ; return (a, cgs_stmts state2) }
getCode :: FCode a -> FCode CmmAGraph
getCode fcode = do { (_,stmts) <- getCodeR fcode; return stmts }
getCodeScoped :: FCode a -> FCode (a, CmmAGraphScoped)
getCodeScoped fcode
  = do  { state1 <- getState
        ; ((a, tscope), state2) <-
            tickScope $
            flip withState state1 { cgs_stmts = mkNop } $
            do { a   <- fcode
               ; scp <- getTickScope
               ; return (a, scp) }
        ; setState $ state2 { cgs_stmts = cgs_stmts state1  }
        ; return (a, (cgs_stmts state2, tscope)) }
getHeapUsage :: (VirtualHpOffset -> FCode a) -> FCode a
getHeapUsage fcode
  = do  { info_down <- getInfoDown
        ; state <- getState
        ; let   fstate_in = state { cgs_hp_usg  = initHpUsage }
                (r, fstate_out) = doFCode (fcode hp_hw) info_down fstate_in
                hp_hw = heapHWM (cgs_hp_usg fstate_out)        
        ; setState $ fstate_out { cgs_hp_usg = cgs_hp_usg state }
        ; return r }
emitCgStmt :: CgStmt -> FCode ()
emitCgStmt stmt
  = do  { state <- getState
        ; setState $ state { cgs_stmts = cgs_stmts state `snocOL` stmt }
        }
emitLabel :: BlockId -> FCode ()
emitLabel id = do tscope <- getTickScope
                  emitCgStmt (CgLabel id tscope)
emitComment :: FastString -> FCode ()
#if 0 /* def DEBUG */
emitComment s = emitCgStmt (CgStmt (CmmComment s))
#else
emitComment _ = return ()
#endif
emitTick :: CmmTickish -> FCode ()
emitTick = emitCgStmt . CgStmt . CmmTick
emitUnwind :: [(GlobalReg, Maybe CmmExpr)] -> FCode ()
emitUnwind regs = do
  dflags <- getDynFlags
  when (debugLevel dflags > 0) $ do
     emitCgStmt $ CgStmt $ CmmUnwind regs
emitAssign :: CmmReg  -> CmmExpr -> FCode ()
emitAssign l r = emitCgStmt (CgStmt (CmmAssign l r))
emitStore :: CmmExpr  -> CmmExpr -> FCode ()
emitStore l r = emitCgStmt (CgStmt (CmmStore l r))
emit :: CmmAGraph -> FCode ()
emit ag
  = do  { state <- getState
        ; setState $ state { cgs_stmts = cgs_stmts state MkGraph.<*> ag } }
emitDecl :: CmmDecl -> FCode ()
emitDecl decl
  = do  { state <- getState
        ; setState $ state { cgs_tops = cgs_tops state `snocOL` decl } }
emitOutOfLine :: BlockId -> CmmAGraphScoped -> FCode ()
emitOutOfLine l (stmts, tscope) = emitCgStmt (CgFork l stmts tscope)
emitProcWithStackFrame
   :: Convention                        
   -> Maybe CmmInfoTable                
   -> CLabel                            
   -> [CmmFormal]                       
   -> [CmmFormal]                       
   -> CmmAGraphScoped                   
   -> Bool                              
   -> FCode ()
emitProcWithStackFrame _conv mb_info lbl _stk_args [] blocks False
  = do  { dflags <- getDynFlags
        ; emitProc_ mb_info lbl [] blocks (widthInBytes (wordWidth dflags)) False
        }
emitProcWithStackFrame conv mb_info lbl stk_args args (graph, tscope) True
        
  = do  { dflags <- getDynFlags
        ; let (offset, live, entry) = mkCallEntry dflags conv args stk_args
              graph' = entry MkGraph.<*> graph
        ; emitProc_ mb_info lbl live (graph', tscope) offset True
        }
emitProcWithStackFrame _ _ _ _ _ _ _ = panic "emitProcWithStackFrame"
emitProcWithConvention :: Convention -> Maybe CmmInfoTable -> CLabel
                       -> [CmmFormal]
                       -> CmmAGraphScoped
                       -> FCode ()
emitProcWithConvention conv mb_info lbl args blocks
  = emitProcWithStackFrame conv mb_info lbl [] args blocks True
emitProc :: Maybe CmmInfoTable -> CLabel -> [GlobalReg] -> CmmAGraphScoped
         -> Int -> FCode ()
emitProc  mb_info lbl live blocks offset
 = emitProc_ mb_info lbl live blocks offset True
emitProc_ :: Maybe CmmInfoTable -> CLabel -> [GlobalReg] -> CmmAGraphScoped
          -> Int -> Bool -> FCode ()
emitProc_ mb_info lbl live blocks offset do_layout
  = do  { dflags <- getDynFlags
        ; l <- newBlockId
        ; let
              blks = labelAGraph l blocks
              infos | Just info <- mb_info = mapSingleton (g_entry blks) info
                    | otherwise            = mapEmpty
              sinfo = StackInfo { arg_space = offset
                                , updfr_space = Just (initUpdFrameOff dflags)
                                , do_layout = do_layout }
              tinfo = TopInfo { info_tbls = infos
                              , stack_info=sinfo}
              proc_block = CmmProc tinfo lbl live blks
        ; state <- getState
        ; setState $ state { cgs_tops = cgs_tops state `snocOL` proc_block } }
getCmm :: FCode () -> FCode CmmGroup
getCmm code
  = do  { state1 <- getState
        ; ((), state2) <- withState code (state1 { cgs_tops  = nilOL })
        ; setState $ state2 { cgs_tops = cgs_tops state1 }
        ; return (fromOL (cgs_tops state2)) }
mkCmmIfThenElse :: CmmExpr -> CmmAGraph -> CmmAGraph -> FCode CmmAGraph
mkCmmIfThenElse e tbranch fbranch = mkCmmIfThenElse' e tbranch fbranch Nothing
mkCmmIfThenElse' :: CmmExpr -> CmmAGraph -> CmmAGraph
                 -> Maybe Bool -> FCode CmmAGraph
mkCmmIfThenElse' e tbranch fbranch likely = do
  tscp  <- getTickScope
  endif <- newBlockId
  tid   <- newBlockId
  fid   <- newBlockId
  let
    (test, then_, else_, likely') = case likely of
      Just False | Just e' <- maybeInvertCmmExpr e
        
        
        
        
        -> (e', fbranch, tbranch, Just True)
      _ -> (e, tbranch, fbranch, likely)
  return $ catAGraphs [ mkCbranch test tid fid likely'
                      , mkLabel tid tscp, then_, mkBranch endif
                      , mkLabel fid tscp, else_, mkLabel endif tscp ]
mkCmmIfGoto :: CmmExpr -> BlockId -> FCode CmmAGraph
mkCmmIfGoto e tid = mkCmmIfGoto' e tid Nothing
mkCmmIfGoto' :: CmmExpr -> BlockId -> Maybe Bool -> FCode CmmAGraph
mkCmmIfGoto' e tid l = do
  endif <- newBlockId
  tscp  <- getTickScope
  return $ catAGraphs [ mkCbranch e tid endif l, mkLabel endif tscp ]
mkCmmIfThen :: CmmExpr -> CmmAGraph -> FCode CmmAGraph
mkCmmIfThen e tbranch = mkCmmIfThen' e tbranch Nothing
mkCmmIfThen' :: CmmExpr -> CmmAGraph -> Maybe Bool -> FCode CmmAGraph
mkCmmIfThen' e tbranch l = do
  endif <- newBlockId
  tid   <- newBlockId
  tscp  <- getTickScope
  return $ catAGraphs [ mkCbranch e tid endif l
                      , mkLabel tid tscp, tbranch, mkLabel endif tscp ]
mkCall :: CmmExpr -> (Convention, Convention) -> [CmmFormal] -> [CmmExpr]
       -> UpdFrameOffset -> [CmmExpr] -> FCode CmmAGraph
mkCall f (callConv, retConv) results actuals updfr_off extra_stack = do
  dflags <- getDynFlags
  k      <- newBlockId
  tscp   <- getTickScope
  let area = Young k
      (off, _, copyin) = copyInOflow dflags retConv area results []
      copyout = mkCallReturnsTo dflags f callConv actuals k off updfr_off extra_stack
  return $ catAGraphs [copyout, mkLabel k tscp, copyin]
mkCmmCall :: CmmExpr -> [CmmFormal] -> [CmmExpr] -> UpdFrameOffset
          -> FCode CmmAGraph
mkCmmCall f results actuals updfr_off
   = mkCall f (NativeDirectCall, NativeReturn) results actuals updfr_off []
aGraphToGraph :: CmmAGraphScoped -> FCode CmmGraph
aGraphToGraph stmts
  = do  { l <- newBlockId
        ; return (labelAGraph l stmts) }