content: {-# 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
  