Main.hs

{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE OverloadedStrings #-}
{-# OPTIONS_GHC -Wall #-}

module Main where

import Control.Monad.Except
import qualified Data.ByteString.Char8 as BS
import Data.IORef
import Data.Int
import qualified Data.Map.Strict as Map
import Foreign.Ptr
import LLVM.AST
import qualified LLVM.AST as AST
import LLVM.AST.Constant
import LLVM.AST.Global
import LLVM.CodeGenOpt
import LLVM.CodeModel
import LLVM.Context
import LLVM.Internal.OrcJIT.CompileLayer
import LLVM.Module
import LLVM.OrcJIT
import LLVM.Relocation
import LLVM.Target
import Prelude hiding (mod)

foreign import ccall "dynamic"
  mkMain :: FunPtr (IO Int32) -> IO Int32

int :: Type
int = IntegerType 32

defAdd :: Definition
defAdd =
  GlobalDefinition
    functionDefaults
      { name = Name "add",
        parameters = ([], False),
        returnType = int,
        basicBlocks = [body]
      }
  where
    body =
      BasicBlock
        (Name "entry")
        []
        (Do $ Ret (Just (ConstantOperand (Int 32 42))) [])

module_ :: AST.Module
module_ =
  defaultModule
    { moduleName = "basic",
      moduleDefinitions = [defAdd]
    }

withTestModule :: AST.Module -> (LLVM.Module.Module -> IO a) -> IO a
withTestModule mod f = withContext $ \context -> withModuleFromAST context mod f

resolver :: CompileLayer l => l -> MangledSymbol -> IO (Either JITSymbolError JITSymbol)
resolver compileLayer symbol = findSymbol compileLayer symbol True

failInIO :: ExceptT String IO a -> IO a
failInIO = either fail return <=< runExceptT

eagerJit :: AST.Module -> IO ()
eagerJit amod = do
  resolvers <- newIORef Map.empty
  withTestModule amod $ \mod ->
    withHostTargetMachine PIC LLVM.CodeModel.Default LLVM.CodeGenOpt.Default $ \tm ->
      withExecutionSession $ \es ->
        withObjectLinkingLayer es (\k -> fmap (\rs -> rs Map.! k) (readIORef resolvers)) $ \linkingLayer ->
          withIRCompileLayer linkingLayer tm $ \compileLayer -> do
            mainSymbol <- mangleSymbol compileLayer "add"
            asm <- moduleLLVMAssembly mod
            BS.putStrLn asm
            withModuleKey es $ \k ->
              withSymbolResolver es (SymbolResolver (resolver compileLayer)) $ \sresolver -> do
                modifyIORef' resolvers (Map.insert k sresolver)
                rsym <- findSymbol compileLayer mainSymbol True
                case rsym of
                  Left err -> do
                    print err
                  Right (JITSymbol mainFn _) -> do
                    result <- mkMain (castPtrToFunPtr (wordPtrToPtr mainFn))
                    print result

main :: IO ()
main = do
  res <- eagerJit module_
  putStrLn "Eager JIT Result:"
  print res