{-# LANGUAGE ForeignFunctionInterface, OverloadedStrings #-}
module LLVM.Test.OrcJIT where

import Test.Tasty
import Test.Tasty.HUnit

import LLVM.Test.Support

import qualified Data.Map.Strict as Map
import Control.Applicative
import Data.ByteString (ByteString)
import Data.Foldable
import Data.Int
import Data.IORef
import Data.Word
import Foreign.Ptr
import Foreign.Storable
import Foreign.Marshal.Alloc
import System.Process (callProcess)
import System.IO.Temp (withSystemTempFile)
import System.IO

import LLVM.Internal.PassManager
import LLVM.Internal.ObjectFile (withObjectFile)
import LLVM.PassManager
import LLVM.Context
import LLVM.Module
import LLVM.OrcJIT
import LLVM.Target
import qualified LLVM.Relocation as Reloc
import qualified LLVM.CodeModel as CodeModel
import qualified LLVM.CodeGenOpt as CodeGenOpt

testModule :: ByteString
testModule =
  "; ModuleID = '<string>'\n\
  \source_filename = \"<string>\"\n\
  \\n\
  \declare i32 @testFunc()\n\
  \define i32 @main(i32, i8**) {\n\
  \  %3 = call i32 @testFunc()\n\
  \  ret i32 %3\n\
  \}\n"

test2Module :: ByteString
test2Module =
  "; ModuleID = '<string>'\n\
  \source_filename = \"<string>\"\n\
  \\n\
  \define i32 @main() {\n\
  \  ret i32 42\n\
  \}\n"

test3Module :: ByteString
test3Module =
  "; ModuleID = '<string>'\n\
  \source_filename = \"<string>\"\n\
  \@data = external global i32\n\
  \\n\
  \define i32 @main() {\n\
  \  %1 = load i32, i32* @data\n\
  \  ret i32 %1\n\
  \}\n"

withTestModule :: ByteString -> (Module -> IO a) -> IO a
withTestModule txt f = withContext $ \context -> withModuleFromLLVMAssembly' context txt f

myTestFuncImpl :: IO Word32
myTestFuncImpl = return 42

foreign import ccall "wrapper"
  wrapTestFunc :: IO Word32 -> IO (FunPtr (IO Word32))

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

tests :: TestTree
tests =
  testGroup "OrcJIT" [
    testCase "basic self-contained function" $ do
      withHostTargetMachine Reloc.PIC CodeModel.Default CodeGenOpt.Default $ \tm ->
        withExecutionSession $ \es -> do
          ol <- createRTDyldObjectLinkingLayer es
          il <- createIRCompileLayer es ol tm
          dylib <- createJITDylib es "testDylib"
          withTestModule test2Module $ \m ->
            withClonedThreadSafeModule m $ \tsm -> do
              addModule tsm dylib il
              Right (JITSymbol addr _) <- lookupSymbol es il dylib "main"
              let mainFn = mkMain (castPtrToFunPtr $ wordPtrToPtr $ fromIntegral addr)
              result <- mainFn
              result @?= 42,

    testCase "basic self-contained function + ObjectLinkingLayer" $ do
      withHostTargetMachine Reloc.PIC CodeModel.Default CodeGenOpt.Default $ \tm ->
        withExecutionSession $ \es -> do
          ol <- createObjectLinkingLayer es
          il <- createIRCompileLayer es ol tm
          dylib <- createJITDylib es "testDylib"
          withTestModule test2Module $ \m ->
            withClonedThreadSafeModule m $ \tsm -> do
              addModule tsm dylib il
              Right (JITSymbol addr _) <- lookupSymbol es il dylib "main"
              let mainFn = mkMain (castPtrToFunPtr $ wordPtrToPtr $ fromIntegral addr)
              result <- mainFn
              result @?= 42,

    testCase "using symbols in external shared libraries" $
      withHostTargetMachine Reloc.PIC CodeModel.Default CodeGenOpt.Default $ \tm ->
        withExecutionSession $ \es -> do
          ol <- createRTDyldObjectLinkingLayer es
          il <- createIRCompileLayer es ol tm
          dylib <- createJITDylib es "testDylib"
          let inputPath = "./test/main_return_38.c"
          withSystemTempFile "main.o" $ \outputPath _ -> do
            callProcess "gcc" ["-shared", "-fPIC", inputPath, "-o", outputPath]
            addDynamicLibrarySearchGenerator il dylib outputPath
            Right (JITSymbol mainFn _) <- lookupSymbol es il dylib "main"
            result <- mkMain (castPtrToFunPtr (wordPtrToPtr mainFn))
            result @?= 38,

    testCase "run optimization passes on a JIT module" $ do
      passmanagerSuccessful <- newIORef False
      withHostTargetMachine Reloc.PIC CodeModel.Default CodeGenOpt.Default $ \tm ->
        withExecutionSession $ \es -> do
          ol <- createRTDyldObjectLinkingLayer es
          il <- createIRCompileLayer es ol tm
          dylib <- createJITDylib es "testDylib"
          withTestModule test2Module $ \m -> do
           withPassManager defaultCuratedPassSetSpec { optLevel = Just 2 } $ \pm -> do
             success <- runPassManager pm m
             writeIORef passmanagerSuccessful success
             withClonedThreadSafeModule m $ \tsm -> do
               addModule tsm dylib il
               Right (JITSymbol mainFn _) <- lookupSymbol es il dylib "main"
               result <- mkMain (castPtrToFunPtr (wordPtrToPtr mainFn))
               result @?= 42
               readIORef passmanagerSuccessful @? "passmanager failed",

    testCase "defining absolute symbols" $ do
      withHostTargetMachine Reloc.PIC CodeModel.Default CodeGenOpt.Default $ \tm ->
        withExecutionSession $ \es -> do
          ol <- createRTDyldObjectLinkingLayer es
          il <- createIRCompileLayer es ol tm
          alloca $ \ptr -> do
            poke ptr (1234 :: Int32)
            dylib <- createJITDylib es "testDylib"
            withMangledSymbol il "data" $ \dataSym -> do
              let flags = defaultJITSymbolFlags { jitSymbolAbsolute = True }
              defineAbsoluteSymbols dylib [(dataSym, JITSymbol (ptrToWordPtr ptr) flags)]
              withTestModule test3Module $ \m -> do
                withClonedThreadSafeModule m $ \tsm -> do
                  addModule tsm dylib il
                  Right (JITSymbol addr _) <- lookupSymbol es il dylib "main"
                  let mainFn = mkMain (castPtrToFunPtr $ wordPtrToPtr $ fromIntegral addr)
                  result <- mainFn
                  result @?= 1234

    -- TODO: Make it possible to use Haskell functions as definition generators
    --       and update to OrcJITv2
    {-
     testCase "eager compilation" $ do
       withTestModule $ \mod ->
         withHostTargetMachine Reloc.PIC CodeModel.Default CodeGenOpt.Default $ \tm ->
         withExecutionSession $ \es ->
         withObjectLinkingLayer es (\k -> fmap (\rs -> rs Map.! k) (readIORef resolvers)) $ \linkingLayer ->
         withIRCompileLayer linkingLayer tm $ \compileLayer -> do
           testFunc <- mangleSymbol compileLayer "testFunc"
           withModuleKey es $ \k ->
             withSymbolResolver es (SymbolResolver (resolver testFunc compileLayer)) $ \resolver -> do
               modifyIORef' resolvers (Map.insert k resolver)
               withModule compileLayer k mod $ do
                 mainSymbol <- mangleSymbol compileLayer "main"
                 Right (JITSymbol mainFn _) <- CL.findSymbol compileLayer mainSymbol True
                 result <- mkMain (castPtrToFunPtr (wordPtrToPtr mainFn))
                 result @?= 42
                 Right (JITSymbol mainFn _) <- CL.findSymbolIn compileLayer k mainSymbol True
                 result <- mkMain (castPtrToFunPtr (wordPtrToPtr mainFn))
                 result @?= 42
                 unknownSymbol <- mangleSymbol compileLayer "unknownSymbol"
                 unknownSymbolRes <- CL.findSymbol compileLayer unknownSymbol True
                 unknownSymbolRes @?= Left (JITSymbolError mempty),
    -}

    -- TODO: Add IRTransformLayer and translate to OrcJITv2
    {-
    testCase "lazy compilation" $ do
      resolvers <- newIORef Map.empty
      let getResolver k = fmap (Map.! k) (readIORef resolvers)
          setResolver k r = modifyIORef' resolvers (Map.insert k r)
      withTestModule $ \mod ->
        withHostTargetMachine Reloc.PIC CodeModel.Default CodeGenOpt.Default $ \tm -> do
          triple <- getTargetMachineTriple tm
          withExecutionSession $ \es ->
            withObjectLinkingLayer es getResolver $ \linkingLayer ->
            withIRCompileLayer linkingLayer tm $ \baseLayer ->
            withIndirectStubsManagerBuilder triple $ \stubsMgr ->
            withJITCompileCallbackManager es triple Nothing $ \callbackMgr ->
            withCompileOnDemandLayer es baseLayer tm getResolver setResolver (\x -> return [x]) callbackMgr stubsMgr False $ \compileLayer -> do
              testFunc <- mangleSymbol compileLayer "testFunc"
              withModuleKey es $ \k ->
                withSymbolResolver es (SymbolResolver (resolver testFunc baseLayer)) $ \resolver -> do
                  setResolver k resolver
                  withModule compileLayer k mod $ do
                    mainSymbol <- mangleSymbol compileLayer "main"
                    Right (JITSymbol mainFn _) <- CL.findSymbol compileLayer mainSymbol True
                    result <- mkMain (castPtrToFunPtr (wordPtrToPtr mainFn))
                    result @?= 42,
    -}

    ]