{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE UndecidableInstances #-}
{- |
Transfer values between Haskell and JIT generated code
in an LLVM-compatible format.
E.g. 'Bool' is stored as 'i1' and Haskell tuples are stored as LLVM structs.
-}
module LLVM.Extra.Marshal where

import qualified LLVM.Extra.Class as Class
import qualified LLVM.Util.Proxy as LP
import qualified LLVM.ExecutionEngine as EE
import qualified LLVM.Core as LLVM
import LLVM.Core (CodeGenFunction, Value)

import qualified Type.Data.Num.Decimal as TypeNum

import Control.Applicative (liftA2, liftA3, (<$>))

import Foreign.Marshal.Alloc (allocaBytes)
import Foreign.Ptr (Ptr)

import Data.Tuple.HT (fst3, snd3, thd3)
import Data.Word (Word8, Word16, Word32, Word64, )
import Data.Int  (Int8,  Int16,  Int32,  Int64, )



peek :: (C a, Struct a ~ struct, EE.Marshal struct) => Ptr struct -> IO a
peek ptr = unpack <$> EE.peek ptr

poke :: (C a, Struct a ~ struct, EE.Marshal struct) => Ptr struct -> a -> IO ()
poke ptr = EE.poke ptr . pack

load ::
   (C a, Struct a ~ struct, EE.Marshal struct) =>
   LP.Proxy a ->
   Value (Ptr struct) -> CodeGenFunction r (Class.ValueTuple a)
load proxy ptr  =  decompose proxy =<< LLVM.load ptr

store ::
   (C a, Struct a ~ struct, EE.Marshal struct) =>
   LP.Proxy a ->
   Class.ValueTuple a -> Value (Ptr struct) -> CodeGenFunction r ()
store proxy tuple ptr  =  flip LLVM.store ptr =<< compose proxy tuple


class
   (Class.MakeValueTuple a, EE.Marshal (Struct a), LLVM.IsSized (Struct a)) =>
      C a where
   type Struct a
   pack :: a -> Struct a
   unpack :: Struct a -> a
   compose ::
      LP.Proxy a ->
      Class.ValueTuple a -> CodeGenFunction r (Value (Struct a))
   decompose ::
      LP.Proxy a ->
      Value (Struct a) -> CodeGenFunction r (Class.ValueTuple a)

instance C Bool where
   type Struct Bool = Bool
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance C Float where
   type Struct Float = Float
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance C Double where
   type Struct Double = Double
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance C Word8 where
   type Struct Word8 = Word8
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance C Word16 where
   type Struct Word16 = Word16
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance C Word32 where
   type Struct Word32 = Word32
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance C Word64 where
   type Struct Word64 = Word64
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance C Int8 where
   type Struct Int8 = Int8
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance C Int16 where
   type Struct Int16 = Int16
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance C Int32 where
   type Struct Int32 = Int32
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance C Int64 where
   type Struct Int64 = Int64
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance (LLVM.IsType a) => C (Ptr a) where
   type Struct (Ptr a) = Ptr a
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance
   (TypeNum.Positive n, TypeNum.Natural (n TypeNum.:*: LLVM.SizeOf a),
    EE.Marshal a, LLVM.IsConst a, LLVM.IsPrimitive a, LLVM.IsSized a) =>
      C (LLVM.Vector n a) where
   type Struct (LLVM.Vector n a) = LLVM.Vector n a
   pack   = id; compose   LP.Proxy = return
   unpack = id; decompose LP.Proxy = return

instance C () where
   type Struct () = LLVM.Struct ()
   pack = LLVM.Struct
   unpack (LLVM.Struct unit) = unit
   compose LP.Proxy () = return $ LLVM.valueOf $ LLVM.Struct ()
   decompose LP.Proxy _ = return ()

instance
   (LLVM.IsSized (Struct a), LLVM.IsSized (Struct b), C a, C b) =>
      C (a,b) where
   type Struct (a,b) = LLVM.Struct (Struct a, (Struct b, ()))
   pack (a,b) = LLVM.Struct (pack a, (pack b, ()))
   unpack (LLVM.Struct (a,(b,()))) = (unpack a, unpack b)
   compose proxy (a,b) = do
      ac <- compose (fst <$> proxy) a
      bc <- compose (snd <$> proxy) b
      struct0 <- LLVM.insertvalue (LLVM.value LLVM.undef) ac TypeNum.d0
      LLVM.insertvalue struct0 bc TypeNum.d1
   decompose proxy struct =
      liftA2 (,)
         (decompose (fst <$> proxy) =<< LLVM.extractvalue struct TypeNum.d0)
         (decompose (snd <$> proxy) =<< LLVM.extractvalue struct TypeNum.d1)

instance
   (LLVM.IsSized (Struct a), LLVM.IsSized (Struct b), LLVM.IsSized (Struct c),
    C a, C b, C c) =>
      C (a,b,c) where
   type Struct (a,b,c) = LLVM.Struct (Struct a, (Struct b, (Struct c, ())))
   pack (a,b,c) = LLVM.Struct (pack a, (pack b, (pack c, ())))
   unpack (LLVM.Struct (a,(b,(c,())))) = (unpack a, unpack b, unpack c)
   compose proxy (a,b,c) = do
      ac <- compose (fst3 <$> proxy) a
      bc <- compose (snd3 <$> proxy) b
      cc <- compose (thd3 <$> proxy) c
      struct0 <- LLVM.insertvalue (LLVM.value LLVM.undef) ac TypeNum.d0
      struct1 <- LLVM.insertvalue struct0 bc TypeNum.d1
      LLVM.insertvalue struct1 cc TypeNum.d2
   decompose proxy struct =
      liftA3 (,,)
         (decompose (fst3 <$> proxy) =<< LLVM.extractvalue struct TypeNum.d0)
         (decompose (snd3 <$> proxy) =<< LLVM.extractvalue struct TypeNum.d1)
         (decompose (thd3 <$> proxy) =<< LLVM.extractvalue struct TypeNum.d2)


with :: (C a) => a -> (Ptr (Struct a) -> IO b) -> IO b
with a act = alloca LP.Proxy $ \ptr -> poke ptr a >> act ptr

alloca ::
   (LLVM.IsType struct) => LP.Proxy struct -> (Ptr struct -> IO b) -> IO b
alloca proxy = allocaBytes (EE.sizeOf proxy)