{-# 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 qualified GHC.Types.Unique.DSM as DSM

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 { unEC :: String -> Env -> Decls -> FCode (Decls, a) }
    deriving (Functor)

type ExtCode = CmmParse ()

returnExtFC :: a -> CmmParse a
returnExtFC a   = EC $ \_ _ s -> return (s, a)

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

instance Applicative CmmParse where
      pure = returnExtFC
      (<*>) = ap

instance Monad CmmParse where
  (>>=) = thenExtFC

instance MonadUnique CmmParse where
  getUniqueSupplyM = code getUniqueSupplyM
  getUniqueM = EC $ \_ _ decls -> do
    u <- getUniqueM
    return (decls, u)

instance DSM.MonadGetUnique CmmParse where
  getUniqueM = GHC.Types.Unique.Supply.getUniqueM

getProfile :: CmmParse Profile
getProfile = EC (\_ _ d -> (d,) <$> F.getProfile)

getPlatform :: CmmParse Platform
getPlatform = EC (\_ _ d -> (d,) <$> F.getPlatform)

getContext :: CmmParse SDocContext
getContext = EC (\_ _ d -> (d,) <$> 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 (EC fcode) =
      EC $ \c e globalDecls -> do
        (_, a) <- F.fixC $ \ ~(decls, _) ->
          fcode c (addListToUFM e decls) globalDecls
        return (globalDecls, a)

-- | Get the current environment from the monad.
getEnv :: CmmParse Env
getEnv  = EC $ \_ e s -> return (s, e)

-- | Get the current context name from the monad
getName :: CmmParse String
getName  = EC $ \c _ s -> return (s, c)

-- | Set context name for a sub-parse
withName :: String -> CmmParse a -> CmmParse a
withName c' (EC fcode) = EC $ \_ e s -> fcode c' e s

addDecl :: FastString -> Named -> ExtCode
addDecl name named = EC $ \_ _ s -> return ((name, named) : s, ())


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

-- | Add a new label to the list of local declarations.
addLabel :: FastString -> BlockId -> ExtCode
addLabel name block_id = addDecl name (LabelN 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 ty name = do
   u <- code newUnique
   let reg = LocalReg u ty
   addVarDecl name (CmmReg (CmmLocal reg))
   return reg


-- | Allocate a fresh label.
newLabel :: FastString -> CmmParse BlockId
newLabel name = do
   u <- code newUnique
   addLabel name (mkBlockId u)
   return (mkBlockId u)

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

newFunctionName name pkg = addDecl name (FunN 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 (name, cmmLabel)
   = addVarDecl name (CmmLit (CmmLabel 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 name = do
  env <- getEnv
  return $
     case lookupUFM env name of
        Just (LabelN l) -> l
        _other          -> mkBlockId (newTagUnique (getUnique name) '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 name = do
  env    <- getEnv
  return $
     case lookupUFM env name of
        Just (VarN e)   -> e
        Just (FunN uid) -> CmmLit (CmmLabel (mkCmmCodeLabel uid       name))
        _other          -> CmmLit (CmmLabel (mkCmmCodeLabel rtsUnitId name))


-- | Lift an FCode computation into the CmmParse monad
code :: FCode a -> CmmParse a
code fc = EC $ \_ _ s -> do
                r <- fc
                return (s, r)

emit :: CmmAGraph -> CmmParse ()
emit = code . F.emit

emitLabel :: BlockId -> CmmParse ()
emitLabel = code . F.emitLabel

emitAssign :: CmmReg  -> CmmExpr -> CmmParse ()
emitAssign l r = code (F.emitAssign l r)

emitStore :: Maybe MemoryOrdering -> CmmExpr  -> CmmExpr -> CmmParse ()
emitStore (Just mem_ord) l r = do
  platform <- getPlatform
  let w = typeWidth $ cmmExprType platform r
  emit $ mkUnsafeCall (PrimTarget $ MO_AtomicWrite w mem_ord) [] [l,r]
emitStore Nothing l r = code (F.emitStore l r)

getCode :: CmmParse a -> CmmParse CmmAGraph
getCode (EC ec) = EC $ \c e s -> do
  ((s',_), gr) <- F.getCodeR (ec c e s)
  return (s', gr)

getCodeR :: CmmParse a -> CmmParse (a, CmmAGraph)
getCodeR (EC ec) = EC $ \c e s -> do
  ((s', r), gr) <- F.getCodeR (ec c e s)
  return (s', (r,gr))

getCodeScoped :: CmmParse a -> CmmParse (a, CmmAGraphScoped)
getCodeScoped (EC ec) = EC $ \c e s -> do
  ((s', r), gr) <- F.getCodeScoped (ec c e s)
  return (s', (r,gr))

emitOutOfLine :: BlockId -> CmmAGraphScoped -> CmmParse ()
emitOutOfLine l g = code (F.emitOutOfLine l g)

withUpdFrameOff :: UpdFrameOffset -> CmmParse () -> CmmParse ()
withUpdFrameOff size inner
  = EC $ \c e s -> F.withUpdFrameOff size $ (unEC inner) c e s

getUpdFrameOff :: CmmParse UpdFrameOffset
getUpdFrameOff = code $ F.getUpdFrameOff