{-# 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 (Ptr, FunPtr, )

import qualified Control.Monad.Trans.Reader as R
import Control.Monad (liftM2, when, void, )
import Control.Applicative (liftA2, pure, (<$>), )

import qualified Data.IORef as IORef
import Data.Functor.Compose (Compose(Compose))

import qualified System.Unsafe as Unsafe

import qualified Synthesizer.LLVM.Debug.Counter as Counter


type Importer f = FunPtr f -> f

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


writeBitcodeToFile :: String -> Counter.T ident -> LLVM.Module -> IO ()
writeBitcodeToFile ext cnt =
   when False .
      LLVM.writeBitcodeToFile
         ("generator" ++ Counter.format 3 cnt ++ ext ++ ".bc")


type Exec = Compose LLVM.CodeGenModule EE.EngineAccess
type Engine = EE.ExecutionEngine

{- |
This function also initializes LLVM.
This simplifies usage from GHCi.
The @llvm@ packages prevents multiple initialization.
-}
compileModule :: Exec externFunction -> IO externFunction
compileModule (Compose bld) = do
   LLVM.initializeNativeTarget
   m <- LLVM.newModule
   (funcs, mappings) <-
      LLVM.defineModule m $ do
         LLVM.setTarget LLVM.hostTriple
         liftM2 (,) bld LLVM.getGlobalMappings

   Counter.with counter $ R.ReaderT $ \cnt -> do
      writeBitcodeToFile "" cnt m
      when False $ do
         void $ Opt.optimizeModule 3 m
         writeBitcodeToFile "-opt" cnt m

   EE.runEngineAccessWithModule m $
      EE.addGlobalMappings mappings >> funcs


createLLVMFunction ::
   (LLVM.FunctionArgs f) =>
   String -> LLVM.FunctionCodeGen f -> LLVM.CodeGenModule (LLVM.Function f)
createLLVMFunction = LLVM.createNamedFunction LLVM.ExternalLinkage

createFunction ::
   (EE.ExecutionFunction f, LLVM.FunctionArgs f) =>
   Importer f -> String -> LLVM.FunctionCodeGen f -> Exec f
createFunction importer name f =
   Compose $ EE.getExecutionFunction importer <$> createLLVMFunction name f


type Finalizer a = (Engine, Ptr a -> IO ())

createFinalizer ::
   (EE.ExecutionFunction f, LLVM.FunctionArgs f) =>
   Importer f -> String -> LLVM.FunctionCodeGen f -> Exec (Engine, f)
createFinalizer importer name f =
   liftA2 (,)
      (Compose $ pure EE.getEngine)
      (createFunction importer name f)