{-# LANGUAGE TypeFamilies #-}
module Synthesizer.LLVM.Execution where

import qualified LLVM.ExecutionEngine as EE
import qualified LLVM.Util.Optimize as Opt
import qualified LLVM.Core as LLVM

import Foreign.Ptr (FunPtr, )

import qualified Control.Monad.Trans.Reader as R
import Control.Monad (liftM2, liftM3, )

import qualified Data.IORef as IORef
import qualified System.Unsafe as Unsafe

import qualified Synthesizer.LLVM.Debug.Counter as Counter


type Importer f = FunPtr f -> f

class Compile externFunction where
   type LLVMFunction externFunction :: *
   compile :: LLVMFunction externFunction -> EE.EngineAccess externFunction

instance Compile (FunPtr f) where
   type LLVMFunction (FunPtr f) = (LLVM.Function f)
   compile = EE.getPointerToFunction

instance (Compile fa, Compile fb) => Compile (fa,fb) where
   type LLVMFunction (fa,fb) = (LLVMFunction fa, LLVMFunction fb)
   compile (fa,fb) =
      liftM2 (,)
         (compile fa)
         (compile fb)

instance (Compile fa, Compile fb, Compile fc) => Compile (fa,fb,fc) where
   type LLVMFunction (fa,fb,fc) = (LLVMFunction fa, LLVMFunction fb, LLVMFunction fc)
   compile (fa,fb,fc) =
      liftM3 (,,)
         (compile fa)
         (compile fb)
         (compile fc)


data BitCodeCnt = BitCodeCnt

{- |
This is only for debugging purposes
and thus I felt free to use unsafePerformIO.
-}
counter :: IORef.IORef (Counter.T BitCodeCnt)
counter =
   Unsafe.performIO $ Counter.new


{- |
This function also initializes LLVM.
This simplifies usage from GHCi.
The @llvm@ packages prevents multiple initialization.
-}
assembleModule ::
   (llvmFunction -> EE.EngineAccess externFunction) ->
   LLVM.CodeGenModule llvmFunction ->
   IO externFunction
assembleModule comp bld = do
   LLVM.initializeNativeTarget
   m <- LLVM.newModule
   (funcs, mappings) <-
      LLVM.defineModule m (liftM2 (,) bld LLVM.getGlobalMappings)

   Counter.with counter $ R.ReaderT $ \cnt -> do
      LLVM.writeBitcodeToFile ("generator" ++ Counter.format 3 cnt ++ ".bc") m
      _ <- Opt.optimizeModule 3 m
      LLVM.writeBitcodeToFile ("generator" ++ Counter.format 3 cnt ++ "-opt.bc") m

   EE.runEngineAccess $
      EE.addModule m >>
      EE.addGlobalMappings mappings >>
      comp funcs


-- this compiles once and is much faster than runFunction
compileModule ::
   (Compile externFunction) =>
   LLVM.CodeGenModule (LLVMFunction externFunction) ->
   IO externFunction
compileModule =
   assembleModule compile

runFunction ::
   (EE.Translatable f) =>
   LLVM.CodeGenModule (LLVM.Function f) -> IO f
runFunction =
   assembleModule EE.generateFunction