{-# LANGUAGE ForeignFunctionInterface #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE DeriveDataTypeable #-}
module LLVM.ExecutionEngine.Engine(
       EngineAccess,
       runEngineAccess,
       createExecutionEngine, addModuleProvider, addModule,
       getExecutionEngineTargetData,
       getPointerToFunction,
       addFunctionValue, addGlobalMappings,
       getFreePointers, FreePointers,
       runFunction, getRunFunction,
       GenericValue, Generic(..)
       ) where

import qualified LLVM.Util.Proxy as Proxy
import qualified LLVM.Core.Util as U

import LLVM.Core.CodeGen (Value(..), Function)
import LLVM.Core.CodeGenMonad (GlobalMappings(..))
import LLVM.Core.Util
          (Module, ModuleProvider, withModuleProvider,
           createModule, createModuleProviderForExistingModule)
import LLVM.Core.Type (IsFirstClass, typeRef)
import LLVM.Util.Proxy (Proxy(Proxy))

import qualified LLVM.FFI.ExecutionEngine as FFI
import qualified LLVM.FFI.Target as FFI
import qualified LLVM.FFI.Core as FFI(ModuleProviderRef, ValueRef)

import qualified Control.Monad.Trans.State as MS
import Control.Monad.IO.Class (MonadIO, liftIO, )
import Control.Monad (liftM, )
import Control.Applicative (Applicative, )
import Control.Concurrent.MVar (MVar, newMVar, putMVar, takeMVar, )

import Data.Typeable (Typeable)
import Data.Int (Int8, Int16, Int32, Int64)
import Data.Word (Word8, Word16, Word32, Word64)

import Foreign.Marshal.Alloc (alloca, free)
import Foreign.Marshal.Array (withArrayLen)
import Foreign.ForeignPtr (ForeignPtr, newForeignPtr, withForeignPtr)
import Foreign.Marshal.Utils (fromBool)
import Foreign.C.String (peekCString)
import Foreign.Ptr (Ptr, FunPtr, )
import Foreign.Storable (peek)
import Foreign.StablePtr (StablePtr, castStablePtrToPtr, castPtrToStablePtr, )
import System.IO.Unsafe (unsafePerformIO)


-- This global variable holds the one and only execution engine.
-- It may be missing, but it never dies.
-- XXX We could provide a destructor, what about functions obtained by runFunction?
{-# NOINLINE theEngine #-}
theEngine :: MVar (Maybe FFI.ExecutionEngineRef)
theEngine = unsafePerformIO $ newMVar Nothing

createExecutionEngine :: ModuleProvider -> IO FFI.ExecutionEngineRef
createExecutionEngine prov =
    withModuleProvider prov $ \provPtr ->
      alloca $ \eePtr ->
        alloca $ \errPtr -> do
          ret <- FFI.createExecutionEngine eePtr provPtr errPtr
          if ret == 1
            then do
                err <- peek errPtr
                errStr <- peekCString err
                free err
                ioError . userError $ errStr
            else
                peek eePtr

getTheEngine :: IO FFI.ExecutionEngineRef
getTheEngine = do
    mee <- takeMVar theEngine
    case mee of
        Just ee -> do putMVar theEngine mee; return ee
        Nothing -> do
            m <- createModule "__empty__"
            mp <- createModuleProviderForExistingModule m
            ee <- createExecutionEngine mp
            putMVar theEngine (Just ee)
            return ee

data EAState = EAState {
    ea_engine :: FFI.ExecutionEngineRef,
    ea_providers :: [ModuleProvider]
    }
    deriving (Show, Typeable)

newtype EngineAccess a = EA (MS.StateT EAState IO a)
    deriving (Functor, Applicative, Monad, MonadIO)

-- |The LLVM execution engine is encapsulated so it cannot be accessed directly.
-- The reason is that (currently) there must only ever be one engine,
-- so access to it is wrapped in a monad.
runEngineAccess :: EngineAccess a -> IO a
runEngineAccess (EA body) = do
    eePtr <- getTheEngine
    MS.evalStateT body $ EAState { ea_engine = eePtr, ea_providers = [] }
    -- XXX should remove module providers again

addModuleProvider :: ModuleProvider -> EngineAccess ()
addModuleProvider prov = do
    ea <- EA MS.get
    EA $ MS.put ea{ ea_providers = prov : ea_providers ea }
    liftIO $ withModuleProvider prov $ \ provPtr ->
                 FFI.addModuleProvider (ea_engine ea) provPtr


getEngine :: EngineAccess FFI.ExecutionEngineRef
getEngine = EA $ MS.gets ea_engine

getExecutionEngineTargetData :: EngineAccess FFI.TargetDataRef
getExecutionEngineTargetData = do
    eePtr <- getEngine
    liftIO $ FFI.getExecutionEngineTargetData eePtr

{- |
In contrast to 'generateFunction' this compiles a function once.
Thus it is faster for many calls to the same function.
See @examples\/Vector.hs@.

If the function calls back into Haskell code,
you also have to set the function addresses
using 'addFunctionValue' or 'addGlobalMappings'.
-}
getPointerToFunction :: Function f -> EngineAccess (FunPtr f)
getPointerToFunction (Value f) = do
    eePtr <- getEngine
    liftIO $ FFI.getPointerToGlobal eePtr f

{- |
Tell LLVM the address of an external function
if it cannot resolve a name automatically.
Alternatively you may declare the function
with 'staticFunction' instead of 'externFunction'.
-}
addFunctionValue :: Function f -> FunPtr f -> EngineAccess ()
addFunctionValue (Value g) f = do
    eePtr <- getEngine
    liftIO $ FFI.addFunctionMapping eePtr g f

{- |
Pass a list of global mappings to LLVM
that can be obtained from 'LLVM.Core.getGlobalMappings'.
-}
addGlobalMappings :: GlobalMappings -> EngineAccess ()
addGlobalMappings (GlobalMappings gms) =
    liftIO . gms =<< getEngine

addModule :: Module -> EngineAccess ()
addModule m = do
    eePtr <- getEngine
    liftIO $ U.withModule m $ FFI.addModule eePtr

-- | Get all the information needed to free a function.
-- Freeing code might have to be done from a (C) finalizer, so it has to done from C.
-- The function c_freeFunctionObject take these pointers as arguments and frees the function.
type FreePointers = (FFI.ExecutionEngineRef, FFI.ModuleProviderRef, FFI.ValueRef)
{-# WARNING getFreePointers "Function returns undefined ModuleProviderRef if there is no module provider" #-}
getFreePointers :: Function f -> EngineAccess FreePointers
getFreePointers (Value f) = do
    ea <- EA MS.get
    liftIO $ do
        let ret mpp = return (ea_engine ea, mpp, f)
        case ea_providers ea of
            prov : _ -> withModuleProvider prov ret
            [] -> ret $ error "getFreePointers: no module provider"

--------------------------------------

newtype GenericValue = GenericValue {
      fromGenericValue :: ForeignPtr FFI.GenericValue
    }

withGenericValue :: GenericValue -> (FFI.GenericValueRef -> IO a) -> IO a
withGenericValue = withForeignPtr . fromGenericValue

createGenericValueWith :: IO FFI.GenericValueRef -> IO GenericValue
createGenericValueWith f = do
  ptr <- f
  liftM GenericValue $ newForeignPtr FFI.ptrDisposeGenericValue ptr

withAll :: [GenericValue] -> (Int -> Ptr FFI.GenericValueRef -> IO a) -> IO a
withAll ps a = go [] ps
    where go ptrs (x:xs) = withGenericValue x $ \ptr -> go (ptr:ptrs) xs
          go ptrs _ = withArrayLen (reverse ptrs) a
                   
runFunction :: U.Function -> [GenericValue] -> EngineAccess GenericValue
runFunction func args = do
    eePtr <- getEngine
    liftIO $ withAll args $ \argLen argPtr ->
                 createGenericValueWith $ FFI.runFunction eePtr func
                                              (fromIntegral argLen) argPtr
getRunFunction :: EngineAccess (U.Function -> [GenericValue] -> IO GenericValue)
getRunFunction = do
    eePtr <- getEngine
    return $ \ func args -> 
             withAll args $ \argLen argPtr ->
                 createGenericValueWith $ FFI.runFunction eePtr func
                                              (fromIntegral argLen) argPtr

class Generic a where
    toGeneric :: a -> GenericValue
    fromGeneric :: GenericValue -> a

instance Generic () where
    toGeneric _ = error "toGeneric ()"
    fromGeneric _ = ()

toGenericInt :: (Integral a, IsFirstClass a) => Bool -> a -> GenericValue
toGenericInt signed val = unsafePerformIO $ createGenericValueWith $ do
    typ <- typeRef $ Proxy.fromValue val
    FFI.createGenericValueOfInt
        typ (fromIntegral val) (fromBool signed)

fromGenericInt :: (Integral a, IsFirstClass a) => Bool -> GenericValue -> a
fromGenericInt signed val = unsafePerformIO $
    withGenericValue val $ \ref ->
        fmap fromIntegral $ FFI.genericValueToInt ref (fromBool signed)

--instance Generic Bool where
--    toGeneric = toGenericInt False . fromBool
--    fromGeneric = toBool . fromGenericInt False

instance Generic Int8 where
    toGeneric = toGenericInt True
    fromGeneric = fromGenericInt True

instance Generic Int16 where
    toGeneric = toGenericInt True
    fromGeneric = fromGenericInt True

instance Generic Int32 where
    toGeneric = toGenericInt True
    fromGeneric = fromGenericInt True

{-
instance Generic Int where
    toGeneric = toGenericInt True
    fromGeneric = fromGenericInt True
-}

instance Generic Int64 where
    toGeneric = toGenericInt True
    fromGeneric = fromGenericInt True

instance Generic Word8 where
    toGeneric = toGenericInt False
    fromGeneric = fromGenericInt False

instance Generic Word16 where
    toGeneric = toGenericInt False
    fromGeneric = fromGenericInt False

instance Generic Word32 where
    toGeneric = toGenericInt False
    fromGeneric = fromGenericInt False

instance Generic Word64 where
    toGeneric = toGenericInt False
    fromGeneric = fromGenericInt False

toGenericReal :: (Real a, IsFirstClass a) => a -> GenericValue
toGenericReal val = unsafePerformIO $ createGenericValueWith $ do
    typ <- typeRef $ Proxy.fromValue val
    FFI.createGenericValueOfFloat typ (realToFrac val)

fromGenericReal :: forall a . (Fractional a, IsFirstClass a) => GenericValue -> a
fromGenericReal val = unsafePerformIO $
    withGenericValue val $ \ ref -> do
        typ <- typeRef (Proxy :: Proxy a)
        fmap realToFrac $ FFI.genericValueToFloat typ ref

instance Generic Float where
    toGeneric = toGenericReal
    fromGeneric = fromGenericReal

instance Generic Double where
    toGeneric = toGenericReal
    fromGeneric = fromGenericReal

instance Generic (Ptr a) where
    toGeneric = unsafePerformIO . createGenericValueWith . FFI.createGenericValueOfPointer
    fromGeneric val = unsafePerformIO . withGenericValue val $ FFI.genericValueToPointer

instance Generic (StablePtr a) where
    toGeneric = unsafePerformIO . createGenericValueWith . FFI.createGenericValueOfPointer . castStablePtrToPtr
    fromGeneric val = unsafePerformIO . fmap castPtrToStablePtr . withGenericValue val $ FFI.genericValueToPointer