ModuleBuilder.hs

{-# LANGUAGE TypeFamilies, MultiParamTypeClasses, UndecidableInstances, TypeOperators #-}

module LLVM.Codegen.ModuleBuilder
  ( ModuleBuilderT
  , ModuleBuilder
  , runModuleBuilderT
  , runModuleBuilder
  , MonadModuleBuilder
  , Module(..)
  , Definition(..)
  , ParameterName(..)
  , FunctionAttribute(..)
  , function
  , global
  , globalUtf8StringPtr
  , extern
  , typedef
  , opaqueTypedef
  , getTypedefs
  , lookupType
  , withFunctionAttributes
  , renderModule
  ) where

import GHC.Stack
import Control.Monad.State.Lazy (StateT(..), MonadState, State, execStateT, modify, gets)
import qualified Control.Monad.State.Strict as StrictState
import qualified Control.Monad.State.Lazy as LazyState
import qualified Control.Monad.RWS.Lazy as LazyRWS
import qualified Control.Monad.RWS.Strict as StrictRWS
import Control.Monad.Reader
import Control.Monad.Writer
import Control.Monad.Except
import Control.Monad.Morph
import Control.Monad.Fix
import Data.DList (DList)
import Data.Map (Map)
import Data.String
import qualified Data.DList as DList
import qualified Data.Map as Map
import qualified Data.Text as T
import qualified Data.Text.Encoding as TE
import qualified Data.ByteString as BS
import Data.Functor.Identity
import LLVM.Codegen.IRBuilder.Monad
import LLVM.Codegen.Operand
import LLVM.Codegen.Type
import LLVM.Codegen.Name
import LLVM.Codegen.Flag
import LLVM.Codegen.IR
import LLVM.Pretty


newtype Module
  = Module [Definition]

data ParameterName
  = ParameterName !T.Text
  | NoParameterName
  deriving Show

instance IsString ParameterName where
  fromString = ParameterName . fromString

data FunctionAttribute
  = WasmExportName !T.Text
  | AlwaysInline
  -- Add more as needed..
  deriving Show

data Global
  = GlobalVariable !Name !Type !Constant
  | Function !Name !Type ![(Type, ParameterName)] ![FunctionAttribute] ![BasicBlock]
  deriving Show

data Typedef
  = Opaque
  | Clear !Type
  deriving Show

data Definition
  = GlobalDefinition !Global
  | TypeDefinition !Name !Typedef
  deriving Show

data ModuleBuilderState
  = ModuleBuilderState
  { definitions :: !(DList Definition)
  , types :: !(Map Name Type)
  , defaultFunctionAttributes :: ![FunctionAttribute]
  }

newtype ModuleBuilderT m a
  = ModuleBuilderT { unModuleBuilderT :: StateT ModuleBuilderState m a }
  deriving ( Functor, Applicative, Monad, MonadFix, MonadIO
           , MonadError e
           )
  via StateT ModuleBuilderState m

type ModuleBuilder = ModuleBuilderT Identity

instance MonadTrans ModuleBuilderT where
  lift = ModuleBuilderT . lift
  {-# INLINEABLE lift #-}

instance MonadReader r m => MonadReader r (ModuleBuilderT m) where
  ask = lift ask
  {-# INLINEABLE ask #-}
  local = mapModuleBuilderT . local
  {-# INLINEABLE local #-}

mapModuleBuilderT :: (Functor m, Monad n) => (m a -> n a) -> ModuleBuilderT m a -> ModuleBuilderT n a
mapModuleBuilderT f (ModuleBuilderT inner) =
  ModuleBuilderT $ do
    s <- LazyState.get
    LazyState.mapStateT (g s) inner
  where
    g s = fmap (,s) . f . fmap fst
{-# INLINEABLE mapModuleBuilderT #-}

instance MonadState s m => MonadState s (ModuleBuilderT m) where
  state = lift . LazyState.state
  {-# INLINEABLE state #-}

instance MFunctor ModuleBuilderT where
  hoist nat = ModuleBuilderT . hoist nat . unModuleBuilderT
  {-# INLINEABLE hoist #-}

class Monad m => MonadModuleBuilder m where
  liftModuleBuilderState :: State ModuleBuilderState a -> m a

  default liftModuleBuilderState
    :: (MonadTrans t, MonadModuleBuilder m1, m ~ t m1)
    => State ModuleBuilderState a
    -> m a
  liftModuleBuilderState = lift . liftModuleBuilderState
  {-# INLINEABLE liftModuleBuilderState #-}

instance Monad m => MonadModuleBuilder (ModuleBuilderT m) where
  liftModuleBuilderState (StateT s) =
    ModuleBuilderT $ StateT $ pure . runIdentity . s
  {-# INLINEABLE liftModuleBuilderState #-}

instance MonadModuleBuilder m => MonadModuleBuilder (IRBuilderT m)
instance MonadModuleBuilder m => MonadModuleBuilder (StrictState.StateT s m)
instance MonadModuleBuilder m => MonadModuleBuilder (LazyState.StateT s m)
instance (MonadModuleBuilder m, Monoid w) => MonadModuleBuilder (StrictRWS.RWST r w s m)
instance (MonadModuleBuilder m, Monoid w) => MonadModuleBuilder (LazyRWS.RWST r w s m)
instance MonadModuleBuilder m => MonadModuleBuilder (ReaderT r m)
instance (MonadModuleBuilder m, Monoid w) => MonadModuleBuilder (WriterT w m)
instance MonadModuleBuilder m => MonadModuleBuilder (ExceptT e m)

runModuleBuilderT :: Monad m => ModuleBuilderT m a -> m Module
runModuleBuilderT (ModuleBuilderT m) =
  Module . DList.toList . definitions <$> execStateT m beginState
  where
    beginState = ModuleBuilderState mempty mempty []
{-# INLINEABLE runModuleBuilderT #-}

withFunctionAttributes
  :: MonadModuleBuilder m
  => ([FunctionAttribute] -> [FunctionAttribute])
  -> m a -> m a
withFunctionAttributes f m = do
  fnAttrs <- liftModuleBuilderState (gets defaultFunctionAttributes)
  liftModuleBuilderState $
    modify $ \s -> s { defaultFunctionAttributes = f fnAttrs }
  result <- m
  liftModuleBuilderState $
    modify $ \s -> s { defaultFunctionAttributes = fnAttrs }
  pure result
{-# INLINEABLE withFunctionAttributes #-}

resetFunctionAttributes :: MonadModuleBuilder m => m ()
resetFunctionAttributes =
  liftModuleBuilderState $
    modify $ \s -> s { defaultFunctionAttributes = mempty }
{-# INLINEABLE resetFunctionAttributes #-}

getDefaultFunctionAttributes :: MonadModuleBuilder m => m [FunctionAttribute]
getDefaultFunctionAttributes =
  liftModuleBuilderState $ gets defaultFunctionAttributes
{-# INLINEABLE getDefaultFunctionAttributes #-}

runModuleBuilder :: ModuleBuilder a -> Module
runModuleBuilder = runIdentity . runModuleBuilderT
{-# INLINEABLE runModuleBuilder #-}

function :: (HasCallStack, MonadModuleBuilder m)
         => Name -> [(Type, ParameterName)] -> Type -> ([Operand] -> IRBuilderT m a) -> m Operand
function name args retTy fnBody = do
  fnAttrs <- getDefaultFunctionAttributes

  (names, instrs) <- runIRBuilderT $ do
    (names, operands) <- unzip <$> traverse (uncurry mkOperand) args
    resetFunctionAttributes  -- This is done to avoid functions emitted in the body that not automatically copy the same attributes
    _ <- fnBody operands
    pure names

  liftModuleBuilderState $
    modify $ \s -> s { defaultFunctionAttributes = fnAttrs }
  let args' = zipWith (\argName (ty, _) -> (ty, ParameterName $ unName argName)) names args
  emitDefinition $ GlobalDefinition $ Function name retTy args' fnAttrs instrs
  pure $ ConstantOperand $ GlobalRef (ptr (FunctionType retTy $ map fst args)) name
{-# INLINEABLE function #-}

emitDefinition :: MonadModuleBuilder m => Definition -> m ()
emitDefinition def =
  liftModuleBuilderState $ modify $ \s -> s { definitions = DList.snoc (definitions s) def }
{-# INLINEABLE emitDefinition #-}

getTypedefs :: MonadModuleBuilder m => m (Map Name Type)
getTypedefs =
  liftModuleBuilderState $ gets types
{-# INLINEABLE getTypedefs #-}

lookupType :: MonadModuleBuilder m => Name -> m (Maybe Type)
lookupType name =
  liftModuleBuilderState $ gets (Map.lookup name . types)
{-# INLINEABLE lookupType #-}

addType :: MonadModuleBuilder m => Name -> Type -> m ()
addType name ty =
  liftModuleBuilderState $ modify $ \s -> s { types = Map.insert name ty (types s) }
{-# INLINEABLE addType #-}

global :: MonadModuleBuilder m => Name -> Type -> Constant -> m Operand
global name ty constant = do
  emitDefinition $ GlobalDefinition $ GlobalVariable name ty constant
  pure $ ConstantOperand $ GlobalRef (ptr ty) name
{-# INLINEABLE global #-}

globalUtf8StringPtr :: (HasCallStack, MonadModuleBuilder m, MonadIRBuilder m)
                    => T.Text -> Name -> m Operand
globalUtf8StringPtr txt name = do
  let utf8Bytes = BS.snoc (TE.encodeUtf8 txt) 0  -- 0-terminated UTF8 string
      llvmValues = map (Int 8 . toInteger) $ BS.unpack utf8Bytes
      arrayValue = Array i8 llvmValues
      constant = ConstantOperand arrayValue
      ty = typeOf constant
  -- This definition will end up before the function this is used in
  addr <- global name ty arrayValue
  let instr = GetElementPtr On addr [ ConstantOperand $ Int 32 0
                                    , ConstantOperand $ Int 32 0
                                    ]
  emitInstr (ptr i8) instr
{-# INLINEABLE globalUtf8StringPtr #-}

-- NOTE: typedefs are only allowed for structs, even though clang also allows it
-- for primitive types. This is done to avoid weird inconsistencies with the LLVM JIT
-- (where this is not allowed).
typedef :: MonadModuleBuilder m => Name -> Flag Packed -> [Type] -> m Type
typedef name packed tys = do
  let ty = StructureType packed tys
  emitDefinition $ TypeDefinition name (Clear ty)
  addType name ty
  pure $ NamedTypeReference name
{-# INLINEABLE typedef #-}

opaqueTypedef :: MonadModuleBuilder m => Name -> m Type
opaqueTypedef name = do
  emitDefinition $ TypeDefinition name Opaque
  pure $ NamedTypeReference name
{-# INLINEABLE opaqueTypedef #-}

extern :: MonadModuleBuilder m => Name -> [Type] -> Type -> m Operand
extern name argTys retTy = do
  let args = [(argTy, ParameterName "") | argTy <- argTys]
  fnAttrs <- getDefaultFunctionAttributes
  emitDefinition $ GlobalDefinition $ Function name retTy args fnAttrs []
  let fnTy = ptr $ FunctionType retTy argTys
  pure $ ConstantOperand $ GlobalRef fnTy name
{-# INLINEABLE extern #-}

-- NOTE: Only used internally, this creates an unassigned operand
mkOperand :: Monad m => Type -> ParameterName -> IRBuilderT m (Name, Operand)
mkOperand ty paramName = do
  name <- case paramName of
    NoParameterName -> freshName Nothing
    ParameterName name -> freshName (Just name)
  pure (name, LocalRef ty name)
{-# INLINEABLE mkOperand #-}

renderModule :: Renderer Module
renderModule buf (Module defs) =
  sepBy "\n\n"# buf defs renderDefinition
{-# INLINEABLE renderModule #-}

renderDefinition :: Renderer Definition
renderDefinition buf = \case
  GlobalDefinition g ->
    renderGlobal buf g
  TypeDefinition name typeDef ->
    case typeDef of
      Opaque ->
        (buf |>. '%') `renderName` name |># " = type opaque"#
      Clear ty ->
        ((buf |>. '%') `renderName` name |># " = type "#) `renderType` ty
{-# INLINEABLE renderDefinition #-}

renderGlobal :: Renderer Global
renderGlobal buf = \case
  GlobalVariable name ty constant ->
    (((((buf |>. '@') `renderName` name) |># " = global "#) `renderType` ty) |>. ' ') `renderConstant` constant
  Function name retTy args attrs body
    | null body ->
      hsep (tupled ((((buf |># "declare external ccc "#) `renderType` retTy) |># " @"#) `renderName` name) argTys renderType
        |># (if null attrs then ""# else " "#)) attrs renderFunctionAttr
    | otherwise ->
      vsep (hsep (tupled ((((buf |># "define external ccc "#) `renderType` retTy) |># " @"#) `renderName` name) (zip [0..] args) renderArg |>. ' ') attrs renderFunctionAttr
        |># (if null attrs then "{\n"# else " {\n"#)) body renderBasicBlock |># "\n}"#
    where
      argTys = map fst args
      renderArg :: Renderer (Int, (Type, ParameterName))
      renderArg buf' (i, (argTy, nm)) =
        let localRef = case nm of
              NoParameterName ->
                LocalRef argTy $ Name $ T.pack $ show i
              ParameterName paramName ->
                LocalRef argTy $ Name paramName
         in ((buf' `renderType` argTy) |>. ' ') `renderOperand` localRef
{-# INLINEABLE renderGlobal #-}

renderFunctionAttr :: Renderer FunctionAttribute
renderFunctionAttr buf = \case
  AlwaysInline ->
    buf |># "alwaysinline"#
  WasmExportName name ->
    dquotes
      (dquotes buf (|># "wasm-export-name"#) |>. '=')
      (|> name)
{-# INLINEABLE renderFunctionAttr #-}