{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables #-}
module LLVM.Internal.OrcJIT where

import LLVM.Prelude

import Control.Exception
import Control.Monad.AnyCont
import Control.Monad.IO.Class
import Data.Bits
import Data.IORef
import Foreign.C.Types
import Foreign.C.String
import Foreign.Ptr

import LLVM.Internal.Coding
import LLVM.Internal.Module (Module, readModule)
import LLVM.Internal.Target (TargetMachine(..))

import qualified LLVM.Internal.FFI.LLVMCTypes as FFI
import qualified LLVM.Internal.FFI.ShortByteString as SBS
import qualified LLVM.Internal.FFI.DataLayout as FFI
import qualified LLVM.Internal.FFI.OrcJIT as FFI
import qualified LLVM.Internal.FFI.Target as FFI

--------------------------------------------------------------------------------
-- ExecutionSession
--------------------------------------------------------------------------------

data ExecutionSession = ExecutionSession {
    sessionPtr :: !(Ptr FFI.ExecutionSession),
    sessionCleanups :: !(IORef [IO ()])
  }

-- | Create a new `ExecutionSession`.
createExecutionSession :: IO ExecutionSession
createExecutionSession = ExecutionSession <$> FFI.createExecutionSession <*> newIORef []

-- | Dispose of an `ExecutionSession`. This should be called when the
-- `ExecutionSession` is not needed anymore.
disposeExecutionSession :: ExecutionSession -> IO ()
disposeExecutionSession (ExecutionSession es cleanups) = do
  FFI.endSession es
  sequence_ =<< readIORef cleanups
  FFI.disposeExecutionSession es

-- | `bracket`-style wrapper around `createExecutionSession` and
-- `disposeExecutionSession`.
withExecutionSession :: (ExecutionSession -> IO a) -> IO a
withExecutionSession = bracket createExecutionSession disposeExecutionSession

--------------------------------------------------------------------------------
-- JITSymbol
--------------------------------------------------------------------------------

-- | Contrary to the C++ interface, we do not store the HasError flag
-- here. Instead decoding a JITSymbol produces a sumtype based on
-- whether that flag is set or not.
data JITSymbolFlags =
  JITSymbolFlags {
    jitSymbolWeak :: !Bool -- ^ Is this a weak symbol?
  , jitSymbolCommon :: !Bool -- ^ Is this a common symbol?
  , jitSymbolAbsolute :: !Bool
    -- ^ Is this an absolute symbol? This will cause LLVM to use
    -- absolute relocations for the symbol even in position
    -- independent code.
  , jitSymbolExported :: !Bool -- ^ Is this symbol exported?
  , jitSymbolCallable :: !Bool
  , jitSymbolMaterializationSideEffectsOnly :: !Bool
  }
  deriving (Show, Eq, Ord)

defaultJITSymbolFlags :: JITSymbolFlags
defaultJITSymbolFlags = JITSymbolFlags False False False False False False

data JITSymbol =
  JITSymbol {
    jitSymbolAddress :: !WordPtr, -- ^ The address of the symbol. If
                                  -- you’ve looked up a function, you
                                  -- need to cast this to a 'FunPtr'.
    jitSymbolFlags :: !JITSymbolFlags -- ^ The flags of this symbol.
  }
  deriving (Show, Eq, Ord)

data JITSymbolError = JITSymbolError ShortByteString
  deriving (Show, Eq)

instance Monad m => EncodeM m JITSymbolFlags FFI.JITSymbolFlags where
  encodeM f = return $ foldr1 (.|.) [
      if a f
         then b
         else 0
    | (a,b) <- [
          (jitSymbolWeak, FFI.jitSymbolFlagsWeak),
          (jitSymbolCommon, FFI.jitSymbolFlagsCommon),
          (jitSymbolAbsolute, FFI.jitSymbolFlagsAbsolute),
          (jitSymbolExported, FFI.jitSymbolFlagsExported),
          (jitSymbolCallable, FFI.jitSymbolFlagsCallable),
          (jitSymbolMaterializationSideEffectsOnly, FFI.jitSymbolFlagsMaterializationSideEffectsOnly)
        ]
    ]

instance Monad m => DecodeM m JITSymbolFlags FFI.JITSymbolFlags where
  decodeM f =
    return $ JITSymbolFlags {
      jitSymbolWeak = FFI.jitSymbolFlagsWeak .&. f /= 0,
      jitSymbolCommon = FFI.jitSymbolFlagsCommon .&. f /= 0,
      jitSymbolAbsolute = FFI.jitSymbolFlagsAbsolute .&. f /= 0,
      jitSymbolExported = FFI.jitSymbolFlagsExported .&. f /= 0,
      jitSymbolCallable = FFI.jitSymbolFlagsCallable .&. f /= 0,
      jitSymbolMaterializationSideEffectsOnly = FFI.jitSymbolFlagsMaterializationSideEffectsOnly .&. f /= 0
    }

instance Monad m => EncodeM m MangledSymbol (Ptr FFI.SymbolStringPtr) where
  encodeM (MangledSymbol p) = return p

instance (Monad m, MonadIO m, MonadAnyCont IO m) => EncodeM m JITSymbol (Ptr FFI.JITEvaluatedSymbol) where
  encodeM (JITSymbol addr flags) = do
    flags' <- encodeM flags
    anyContToM $ bracket
      (FFI.createJITEvaluatedSymbol (FFI.TargetAddress $ fromIntegral addr) flags')
      FFI.disposeJITEvaluatedSymbol

instance (MonadIO m, MonadAnyCont IO m) =>
         DecodeM m (Either JITSymbolError JITSymbol) (Ptr FFI.ExpectedJITEvaluatedSymbol) where
  decodeM expectedSym = do
    errMsg <- alloca
    FFI.TargetAddress addr <- liftIO $ FFI.getExpectedSymbolAddress expectedSym errMsg
    rawFlags <- liftIO (FFI.getExpectedSymbolFlags expectedSym)
    if addr == 0 || (rawFlags .&. FFI.jitSymbolFlagsHasError /= 0)
      then do
        errMsg <- decodeM errMsg
        pure (Left (JITSymbolError errMsg))
      else do
        flags <- decodeM rawFlags
        pure (Right (JITSymbol (fromIntegral addr) flags))

--------------------------------------------------------------------------------
-- JITDylib
--------------------------------------------------------------------------------

newtype JITDylib = JITDylib (Ptr FFI.JITDylib)

-- | Create a new 'JITDylib' with the given name.
createJITDylib :: ExecutionSession -> ShortByteString -> IO JITDylib
createJITDylib (ExecutionSession es _) name =
  SBS.useAsCString name $ fmap JITDylib . FFI.createJITDylib es

-- NB: JITDylib unloading is WIP (at least according to some old-looking docs)

-- | Adds a 'JITDylib' definition generator that looks up missing symbols in
-- the namespace of the current process.
addDynamicLibrarySearchGeneratorForCurrentProcess :: IRLayer l => l -> JITDylib -> IO ()
addDynamicLibrarySearchGeneratorForCurrentProcess compileLayer (JITDylib dylib) =
  FFI.addDynamicLibrarySearchGeneratorForCurrentProcess dylib (getDataLayout compileLayer)

-- | Adds a 'JITDylib' definition generator that looks up missing symbols in
-- the namespace of a shared library located at the specified 'FilePath'.
addDynamicLibrarySearchGenerator :: IRLayer l => l -> JITDylib -> FilePath -> IO ()
addDynamicLibrarySearchGenerator compileLayer (JITDylib dylib) s = withCString s $ \cStr ->
  FFI.addDynamicLibrarySearchGenerator dylib (getDataLayout compileLayer) cStr

defineAbsoluteSymbols :: JITDylib -> [(MangledSymbol, JITSymbol)] -> IO ()
defineAbsoluteSymbols (JITDylib dylib) symList =
  runAnyContT' return $ do
    (nsyms, symStrPtrPtr) :: (CUInt, Ptr (Ptr FFI.SymbolStringPtr)) <- encodeM symNames
    (_, symValPtrPtr) :: (CUInt, Ptr (Ptr FFI.JITEvaluatedSymbol)) <- encodeM symVals
    liftIO $ FFI.defineAbsoluteSymbols dylib nsyms symStrPtrPtr symValPtrPtr
  where (symNames, symVals) = unzip symList

-- | Looks up an (unmangled) symbol name in the given 'JITDylib'.
--
-- The symbol is expected to have been added to the 'JITDylib' by the same 'IRLayer'
-- as specified in this function. Using a different 'IRLayer' can cause the lookup
-- to fail due to differences in mangling schemes.
lookupSymbol :: IRLayer l => ExecutionSession -> l -> JITDylib -> ShortByteString -> IO (Either JITSymbolError JITSymbol)
lookupSymbol (ExecutionSession es _) irl (JITDylib dylib) name = SBS.useAsCString name $ \nameStr ->
  runAnyContT' return $ do
    symbol <- anyContToM $ bracket
      (FFI.lookupSymbol es dylib (getMangler irl) nameStr) FFI.disposeExpectedJITEvaluatedSymbol
    decodeM symbol

--------------------------------------------------------------------------------
-- ThreadSafeContext
--------------------------------------------------------------------------------

newtype ThreadSafeContext = ThreadSafeContext (Ptr FFI.ThreadSafeContext)

-- | Create a 'ThreadSafeContext'
createThreadSafeContext :: IO ThreadSafeContext
createThreadSafeContext = ThreadSafeContext <$> FFI.createThreadSafeContext

-- | Dispose of a 'ThreadSafeContext'
disposeThreadSafeContext :: ThreadSafeContext -> IO ()
disposeThreadSafeContext (ThreadSafeContext ctx) = FFI.disposeThreadSafeContext ctx

-- | 'bracket'-style wrapper around 'createThreadSafeContext'
-- and 'disposeThreadSafeContext'.
withThreadSafeContext :: (ThreadSafeContext -> IO a) -> IO a
withThreadSafeContext = bracket createThreadSafeContext disposeThreadSafeContext

--------------------------------------------------------------------------------
-- ThreadSafeModule
--------------------------------------------------------------------------------

newtype ThreadSafeModule = ThreadSafeModule (Ptr FFI.ThreadSafeModule)

-- | Create a 'ThreadSafeModule' with the same content as the input 'Module'.
--
-- The module will get cloned into a fresh LLVM context. The lifetime of the
-- new context is bound to the lifetime of the returned 'ThreadSafeModule'.
cloneAsThreadSafeModule :: Module -> IO ThreadSafeModule
cloneAsThreadSafeModule m = do
  mPtr <- readModule m
  ThreadSafeModule <$> FFI.cloneAsThreadSafeModule mPtr

-- | Dispose of a 'ThreadSafeModule'.
disposeThreadSafeModule :: ThreadSafeModule -> IO ()
disposeThreadSafeModule (ThreadSafeModule m) = FFI.disposeThreadSafeModule m

-- | 'bracket'-style wrapper around 'cloneAsThreadSafeModule'
-- and 'disposeThreadSafeModule'.
withClonedThreadSafeModule :: Module -> (ThreadSafeModule -> IO a) -> IO a
withClonedThreadSafeModule m = bracket (cloneAsThreadSafeModule m) disposeThreadSafeModule

--------------------------------------------------------------------------------
-- ObjectLayer + RTDyldObjectLinkingLayer
--------------------------------------------------------------------------------

-- | A type class implemented by the different OrcJIT object layers.
--
-- See e.g. 'RTDyldObjectLinkingLayer'.
class ObjectLayer l where
  getObjectLayer :: l -> Ptr FFI.ObjectLayer


data RTDyldObjectLinkingLayer = RTDyldObjectLinkingLayer !(Ptr FFI.ObjectLayer)

instance ObjectLayer RTDyldObjectLinkingLayer where
  getObjectLayer (RTDyldObjectLinkingLayer ol) = ol

-- | Create a new 'RTDyldObjectLinkingLayer'.
--
-- The layer will get automatically disposed along with its ExecutionSession.
createRTDyldObjectLinkingLayer :: ExecutionSession -> IO RTDyldObjectLinkingLayer
createRTDyldObjectLinkingLayer (ExecutionSession es cleanups) = do
  ol <- FFI.createRTDyldObjectLinkingLayer es
  modifyIORef' cleanups (FFI.disposeObjectLayer ol :)
  return $ RTDyldObjectLinkingLayer ol

data ObjectLinkingLayer = ObjectLinkingLayer !(Ptr FFI.ObjectLayer)

instance ObjectLayer ObjectLinkingLayer where
  getObjectLayer (ObjectLinkingLayer ol) = ol

createObjectLinkingLayer :: ExecutionSession -> IO ObjectLinkingLayer
createObjectLinkingLayer (ExecutionSession es cleanups) = do
  ol <- FFI.createObjectLinkingLayer es
  modifyIORef' cleanups (FFI.disposeObjectLayer ol :)
  return $ ObjectLinkingLayer ol

addObjectFile :: ObjectLayer l => l -> JITDylib -> FilePath -> IO ()
addObjectFile ol (JITDylib dylib) path = do
  withCString path $ \cStr ->
    FFI.objectLayerAddObjectFile (getObjectLayer ol) dylib cStr

--------------------------------------------------------------------------------
-- IRLayer + IRCompileLayer
--------------------------------------------------------------------------------

-- | A mangled symbol name. Valid only for as long as the IRLayer that created it.
newtype MangledSymbol = MangledSymbol (Ptr FFI.SymbolStringPtr)

-- | A type class implemented by the different OrcJIT IR layers.
--
-- See e.g. 'IRCompileLayer'.
class IRLayer l where
  getIRLayer :: l -> Ptr FFI.IRLayer
  getDataLayout :: l -> Ptr FFI.DataLayout
  getMangler :: l -> Ptr FFI.MangleAndInterner

-- | Add a 'Module' to the specified 'JITDylib'.
--
-- The specified 'IRLayer' will be responsible for compiling the symbols
-- present in the module. The module itself is consumed and __should not be used again__.
addModule :: IRLayer l => ThreadSafeModule -> JITDylib -> l -> IO ()
addModule (ThreadSafeModule m) (JITDylib dylib) irl =
  FFI.irLayerAddModule m dylib (getDataLayout irl) (getIRLayer irl)

mangleSymbol :: IRLayer l => l -> ShortByteString -> IO MangledSymbol
mangleSymbol irl name = SBS.useAsCString name $ \namePtr ->
  MangledSymbol <$> FFI.mangleSymbol (getMangler irl) namePtr

disposeMangledSymbol :: MangledSymbol -> IO ()
disposeMangledSymbol (MangledSymbol symbol) = FFI.disposeMangledSymbol symbol

withMangledSymbol :: IRLayer l => l -> ShortByteString -> (MangledSymbol -> IO a) -> IO a
withMangledSymbol irl name = bracket (mangleSymbol irl name) disposeMangledSymbol

-- | An IR layer that compiles the symbols in a module eagerly.
data IRCompileLayer = IRCompileLayer !(Ptr FFI.IRLayer) !(Ptr FFI.DataLayout) !(Ptr FFI.MangleAndInterner)

instance IRLayer IRCompileLayer where
  getIRLayer    (IRCompileLayer cl _ _) = cl
  getDataLayout (IRCompileLayer _ dl _) = dl
  getMangler    (IRCompileLayer _ _ mg) = mg

-- | Create a new 'IRCompileLayer'.
--
-- The layer will get automatically disposed along with its ExecutionSession.
createIRCompileLayer :: ObjectLayer l => ExecutionSession -> l -> TargetMachine -> IO IRCompileLayer
createIRCompileLayer (ExecutionSession es cleanups) ol (TargetMachine tm) = do
  dl <- FFI.createTargetDataLayout tm
  modifyIORef' cleanups (FFI.disposeDataLayout dl :)
  mg <- FFI.createMangleAndInterner es dl
  modifyIORef' cleanups (FFI.disposeMangleAndInterner mg :)
  cl <- FFI.createIRCompileLayer es (getObjectLayer ol) tm
  modifyIORef' cleanups (FFI.disposeIRLayer cl :)
  return $ IRCompileLayer cl dl mg