{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE TupleSections #-}
-- | Our extended FCode monad.

-- We add a mapping from names to CmmExpr, to support local variable names in
-- the concrete C-- code.  The unique supply of the underlying FCode monad
-- is used to grab a new unique for each local variable.

-- In C--, a local variable can be declared anywhere within a proc,
-- and it scopes from the beginning of the proc to the end.  Hence, we have
-- to collect declarations as we parse the proc, and feed the environment
-- back in circularly (to avoid a two-pass algorithm).

module GHC.StgToCmm.ExtCode (
        CmmParse, unEC,
        Named(..), Env,

        loopDecls,
        getEnv,

        withName,
        getName,

        newLocal,
        newLabel,
        newBlockId,
        newFunctionName,
        newImport,
        lookupLabel,
        lookupName,

        code,
        emit, emitLabel, emitAssign, emitStore,
        getCode, getCodeR, getCodeScoped,
        emitOutOfLine,
        withUpdFrameOff, getUpdFrameOff,
        getProfile, getPlatform, getContext
)

where

import GHC.Prelude

import GHC.Platform
import GHC.Platform.Profile

import qualified GHC.StgToCmm.Monad as F
import GHC.StgToCmm.Monad (FCode, newUnique)

import GHC.Cmm
import GHC.Cmm.CLabel
import GHC.Cmm.Graph

import GHC.Cmm.BlockId
import GHC.Data.FastString
import GHC.Unit.Module
import GHC.Types.Unique.FM
import GHC.Types.Unique
import GHC.Types.Unique.Supply

import Control.Monad (ap)
import GHC.Utils.Outputable (SDocContext)

-- | The environment contains variable definitions or blockids.
data Named
        = VarN CmmExpr          -- ^ Holds CmmLit(CmmLabel ..) which gives the label type,
                                --      eg, RtsLabel, ForeignLabel, CmmLabel etc.

        | FunN   UnitId         -- ^ A function name from this unit
        | LabelN BlockId        -- ^ A blockid of some code or data.

-- | An environment of named things.
type Env        = UniqFM FastString Named

-- | Local declarations that are in scope during code generation.
type Decls      = [(FastString,Named)]

-- | Does a computation in the FCode monad, with a current environment
--      and a list of local declarations. Returns the resulting list of declarations.
newtype CmmParse a
        = EC { forall a. CmmParse a -> String -> Env -> Decls -> FCode (Decls, a)
unEC :: String -> Env -> Decls -> FCode (Decls, a) }
    deriving (forall a b. a -> CmmParse b -> CmmParse a
forall a b. (a -> b) -> CmmParse a -> CmmParse b
forall (f :: * -> *).
(forall a b. (a -> b) -> f a -> f b)
-> (forall a b. a -> f b -> f a) -> Functor f
<$ :: forall a b. a -> CmmParse b -> CmmParse a
$c<$ :: forall a b. a -> CmmParse b -> CmmParse a
fmap :: forall a b. (a -> b) -> CmmParse a -> CmmParse b
$cfmap :: forall a b. (a -> b) -> CmmParse a -> CmmParse b
Functor)

type ExtCode = CmmParse ()

returnExtFC :: a -> CmmParse a
returnExtFC :: forall a. a -> CmmParse a
returnExtFC a
a   = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
_ Env
_ Decls
s -> forall (m :: * -> *) a. Monad m => a -> m a
return (Decls
s, a
a)

thenExtFC :: CmmParse a -> (a -> CmmParse b) -> CmmParse b
thenExtFC :: forall a b. CmmParse a -> (a -> CmmParse b) -> CmmParse b
thenExtFC (EC String -> Env -> Decls -> FCode (Decls, a)
m) a -> CmmParse b
k = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
c Env
e Decls
s -> do (Decls
s',a
r) <- String -> Env -> Decls -> FCode (Decls, a)
m String
c Env
e Decls
s; forall a. CmmParse a -> String -> Env -> Decls -> FCode (Decls, a)
unEC (a -> CmmParse b
k a
r) String
c Env
e Decls
s'

instance Applicative CmmParse where
      pure :: forall a. a -> CmmParse a
pure = forall a. a -> CmmParse a
returnExtFC
      <*> :: forall a b. CmmParse (a -> b) -> CmmParse a -> CmmParse b
(<*>) = forall (m :: * -> *) a b. Monad m => m (a -> b) -> m a -> m b
ap

instance Monad CmmParse where
  >>= :: forall a b. CmmParse a -> (a -> CmmParse b) -> CmmParse b
(>>=) = forall a b. CmmParse a -> (a -> CmmParse b) -> CmmParse b
thenExtFC

instance MonadUnique CmmParse where
  getUniqueSupplyM :: CmmParse UniqSupply
getUniqueSupplyM = forall a. FCode a -> CmmParse a
code forall (m :: * -> *). MonadUnique m => m UniqSupply
getUniqueSupplyM
  getUniqueM :: CmmParse Unique
getUniqueM = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
_ Env
_ Decls
decls -> do
    Unique
u <- forall (m :: * -> *). MonadUnique m => m Unique
getUniqueM
    forall (m :: * -> *) a. Monad m => a -> m a
return (Decls
decls, Unique
u)

getProfile :: CmmParse Profile
getProfile :: CmmParse Profile
getProfile = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC (\String
_ Env
_ Decls
d -> (Decls
d,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FCode Profile
F.getProfile)

getPlatform :: CmmParse Platform
getPlatform :: CmmParse Platform
getPlatform = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC (\String
_ Env
_ Decls
d -> (Decls
d,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FCode Platform
F.getPlatform)

getContext :: CmmParse SDocContext
getContext :: CmmParse SDocContext
getContext = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC (\String
_ Env
_ Decls
d -> (Decls
d,) forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> FCode SDocContext
F.getContext)

-- | Takes the variable declarations and imports from the monad
--      and makes an environment, which is looped back into the computation.
--      In this way, we can have embedded declarations that scope over the whole
--      procedure, and imports that scope over the entire module.
--      Discards the local declaration contained within decl'
--
loopDecls :: CmmParse a -> CmmParse a
loopDecls :: forall a. CmmParse a -> CmmParse a
loopDecls (EC String -> Env -> Decls -> FCode (Decls, a)
fcode) =
      forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
c Env
e Decls
globalDecls -> do
        (Decls
_, a
a) <- forall a. (a -> FCode a) -> FCode a
F.fixC forall a b. (a -> b) -> a -> b
$ \ ~(Decls
decls, a
_) ->
          String -> Env -> Decls -> FCode (Decls, a)
fcode String
c (forall key elt.
Uniquable key =>
UniqFM key elt -> [(key, elt)] -> UniqFM key elt
addListToUFM Env
e Decls
decls) Decls
globalDecls
        forall (m :: * -> *) a. Monad m => a -> m a
return (Decls
globalDecls, a
a)

-- | Get the current environment from the monad.
getEnv :: CmmParse Env
getEnv :: CmmParse Env
getEnv  = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
_ Env
e Decls
s -> forall (m :: * -> *) a. Monad m => a -> m a
return (Decls
s, Env
e)

-- | Get the current context name from the monad
getName :: CmmParse String
getName :: CmmParse String
getName  = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
c Env
_ Decls
s -> forall (m :: * -> *) a. Monad m => a -> m a
return (Decls
s, String
c)

-- | Set context name for a sub-parse
withName :: String -> CmmParse a -> CmmParse a
withName :: forall a. String -> CmmParse a -> CmmParse a
withName String
c' (EC String -> Env -> Decls -> FCode (Decls, a)
fcode) = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
_ Env
e Decls
s -> String -> Env -> Decls -> FCode (Decls, a)
fcode String
c' Env
e Decls
s

addDecl :: FastString -> Named -> ExtCode
addDecl :: FastString -> Named -> ExtCode
addDecl FastString
name Named
named = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
_ Env
_ Decls
s -> forall (m :: * -> *) a. Monad m => a -> m a
return ((FastString
name, Named
named) forall a. a -> [a] -> [a]
: Decls
s, ())


-- | Add a new variable to the list of local declarations.
--      The CmmExpr says where the value is stored.
addVarDecl :: FastString -> CmmExpr -> ExtCode
addVarDecl :: FastString -> CmmExpr -> ExtCode
addVarDecl FastString
var CmmExpr
expr = FastString -> Named -> ExtCode
addDecl FastString
var (CmmExpr -> Named
VarN CmmExpr
expr)

-- | Add a new label to the list of local declarations.
addLabel :: FastString -> BlockId -> ExtCode
addLabel :: FastString -> BlockId -> ExtCode
addLabel FastString
name BlockId
block_id = FastString -> Named -> ExtCode
addDecl FastString
name (BlockId -> Named
LabelN BlockId
block_id)


-- | Create a fresh local variable of a given type.
newLocal
        :: CmmType              -- ^ data type
        -> FastString           -- ^ name of variable
        -> CmmParse LocalReg    -- ^ register holding the value

newLocal :: CmmType -> FastString -> CmmParse LocalReg
newLocal CmmType
ty FastString
name = do
   Unique
u <- forall a. FCode a -> CmmParse a
code FCode Unique
newUnique
   let reg :: LocalReg
reg = Unique -> CmmType -> LocalReg
LocalReg Unique
u CmmType
ty
   FastString -> CmmExpr -> ExtCode
addVarDecl FastString
name (CmmReg -> CmmExpr
CmmReg (LocalReg -> CmmReg
CmmLocal LocalReg
reg))
   forall (m :: * -> *) a. Monad m => a -> m a
return LocalReg
reg


-- | Allocate a fresh label.
newLabel :: FastString -> CmmParse BlockId
newLabel :: FastString -> CmmParse BlockId
newLabel FastString
name = do
   Unique
u <- forall a. FCode a -> CmmParse a
code FCode Unique
newUnique
   FastString -> BlockId -> ExtCode
addLabel FastString
name (Unique -> BlockId
mkBlockId Unique
u)
   forall (m :: * -> *) a. Monad m => a -> m a
return (Unique -> BlockId
mkBlockId Unique
u)

-- | Add a local function to the environment.
newFunctionName
        :: FastString   -- ^ name of the function
        -> UnitId       -- ^ package of the current module
        -> ExtCode

newFunctionName :: FastString -> UnitId -> ExtCode
newFunctionName FastString
name UnitId
pkg = FastString -> Named -> ExtCode
addDecl FastString
name (UnitId -> Named
FunN UnitId
pkg)


-- | Add an imported foreign label to the list of local declarations.
--      If this is done at the start of the module the declaration will scope
--      over the whole module.
newImport
        :: (FastString, CLabel)
        -> CmmParse ()

newImport :: (FastString, CLabel) -> ExtCode
newImport (FastString
name, CLabel
cmmLabel)
   = FastString -> CmmExpr -> ExtCode
addVarDecl FastString
name (CmmLit -> CmmExpr
CmmLit (CLabel -> CmmLit
CmmLabel CLabel
cmmLabel))


-- | Lookup the BlockId bound to the label with this name.
--      If one hasn't been bound yet, create a fresh one based on the
--      Unique of the name.
lookupLabel :: FastString -> CmmParse BlockId
lookupLabel :: FastString -> CmmParse BlockId
lookupLabel FastString
name = do
  Env
env <- CmmParse Env
getEnv
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
     case forall key elt. Uniquable key => UniqFM key elt -> key -> Maybe elt
lookupUFM Env
env FastString
name of
        Just (LabelN BlockId
l) -> BlockId
l
        Maybe Named
_other          -> Unique -> BlockId
mkBlockId (Unique -> Char -> Unique
newTagUnique (forall a. Uniquable a => a -> Unique
getUnique FastString
name) Char
'L')


-- | Lookup the location of a named variable.
--      Unknown names are treated as if they had been 'import'ed from the runtime system.
--      This saves us a lot of bother in the RTS sources, at the expense of
--      deferring some errors to link time.
lookupName :: FastString -> CmmParse CmmExpr
lookupName :: FastString -> CmmParse CmmExpr
lookupName FastString
name = do
  Env
env    <- CmmParse Env
getEnv
  forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$
     case forall key elt. Uniquable key => UniqFM key elt -> key -> Maybe elt
lookupUFM Env
env FastString
name of
        Just (VarN CmmExpr
e)   -> CmmExpr
e
        Just (FunN UnitId
uid) -> CmmLit -> CmmExpr
CmmLit (CLabel -> CmmLit
CmmLabel (UnitId -> FastString -> CLabel
mkCmmCodeLabel UnitId
uid       FastString
name))
        Maybe Named
_other          -> CmmLit -> CmmExpr
CmmLit (CLabel -> CmmLit
CmmLabel (UnitId -> FastString -> CLabel
mkCmmCodeLabel UnitId
rtsUnitId FastString
name))


-- | Lift an FCode computation into the CmmParse monad
code :: FCode a -> CmmParse a
code :: forall a. FCode a -> CmmParse a
code FCode a
fc = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
_ Env
_ Decls
s -> do
                a
r <- FCode a
fc
                forall (m :: * -> *) a. Monad m => a -> m a
return (Decls
s, a
r)

emit :: CmmAGraph -> CmmParse ()
emit :: CmmAGraph -> ExtCode
emit = forall a. FCode a -> CmmParse a
code forall b c a. (b -> c) -> (a -> b) -> a -> c
. CmmAGraph -> FCode ()
F.emit

emitLabel :: BlockId -> CmmParse ()
emitLabel :: BlockId -> ExtCode
emitLabel = forall a. FCode a -> CmmParse a
code forall b c a. (b -> c) -> (a -> b) -> a -> c
. BlockId -> FCode ()
F.emitLabel

emitAssign :: CmmReg  -> CmmExpr -> CmmParse ()
emitAssign :: CmmReg -> CmmExpr -> ExtCode
emitAssign CmmReg
l CmmExpr
r = forall a. FCode a -> CmmParse a
code (CmmReg -> CmmExpr -> FCode ()
F.emitAssign CmmReg
l CmmExpr
r)

emitStore :: CmmExpr  -> CmmExpr -> CmmParse ()
emitStore :: CmmExpr -> CmmExpr -> ExtCode
emitStore CmmExpr
l CmmExpr
r = forall a. FCode a -> CmmParse a
code (CmmExpr -> CmmExpr -> FCode ()
F.emitStore CmmExpr
l CmmExpr
r)

getCode :: CmmParse a -> CmmParse CmmAGraph
getCode :: forall a. CmmParse a -> CmmParse CmmAGraph
getCode (EC String -> Env -> Decls -> FCode (Decls, a)
ec) = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
c Env
e Decls
s -> do
  ((Decls
s',a
_), CmmAGraph
gr) <- forall a. FCode a -> FCode (a, CmmAGraph)
F.getCodeR (String -> Env -> Decls -> FCode (Decls, a)
ec String
c Env
e Decls
s)
  forall (m :: * -> *) a. Monad m => a -> m a
return (Decls
s', CmmAGraph
gr)

getCodeR :: CmmParse a -> CmmParse (a, CmmAGraph)
getCodeR :: forall a. CmmParse a -> CmmParse (a, CmmAGraph)
getCodeR (EC String -> Env -> Decls -> FCode (Decls, a)
ec) = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
c Env
e Decls
s -> do
  ((Decls
s', a
r), CmmAGraph
gr) <- forall a. FCode a -> FCode (a, CmmAGraph)
F.getCodeR (String -> Env -> Decls -> FCode (Decls, a)
ec String
c Env
e Decls
s)
  forall (m :: * -> *) a. Monad m => a -> m a
return (Decls
s', (a
r,CmmAGraph
gr))

getCodeScoped :: CmmParse a -> CmmParse (a, CmmAGraphScoped)
getCodeScoped :: forall a. CmmParse a -> CmmParse (a, CmmAGraphScoped)
getCodeScoped (EC String -> Env -> Decls -> FCode (Decls, a)
ec) = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
c Env
e Decls
s -> do
  ((Decls
s', a
r), CmmAGraphScoped
gr) <- forall a. FCode a -> FCode (a, CmmAGraphScoped)
F.getCodeScoped (String -> Env -> Decls -> FCode (Decls, a)
ec String
c Env
e Decls
s)
  forall (m :: * -> *) a. Monad m => a -> m a
return (Decls
s', (a
r,CmmAGraphScoped
gr))

emitOutOfLine :: BlockId -> CmmAGraphScoped -> CmmParse ()
emitOutOfLine :: BlockId -> CmmAGraphScoped -> ExtCode
emitOutOfLine BlockId
l CmmAGraphScoped
g = forall a. FCode a -> CmmParse a
code (BlockId -> CmmAGraphScoped -> FCode ()
F.emitOutOfLine BlockId
l CmmAGraphScoped
g)

withUpdFrameOff :: UpdFrameOffset -> CmmParse () -> CmmParse ()
withUpdFrameOff :: UpdFrameOffset -> ExtCode -> ExtCode
withUpdFrameOff UpdFrameOffset
size ExtCode
inner
  = forall a.
(String -> Env -> Decls -> FCode (Decls, a)) -> CmmParse a
EC forall a b. (a -> b) -> a -> b
$ \String
c Env
e Decls
s -> forall a. UpdFrameOffset -> FCode a -> FCode a
F.withUpdFrameOff UpdFrameOffset
size forall a b. (a -> b) -> a -> b
$ (forall a. CmmParse a -> String -> Env -> Decls -> FCode (Decls, a)
unEC ExtCode
inner) String
c Env
e Decls
s

getUpdFrameOff :: CmmParse UpdFrameOffset
getUpdFrameOff :: CmmParse UpdFrameOffset
getUpdFrameOff = forall a. FCode a -> CmmParse a
code forall a b. (a -> b) -> a -> b
$ FCode UpdFrameOffset
F.getUpdFrameOff