{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE FlexibleContexts #-}
module LLVM.Extra.Memory (
   C(load, store, decompose, compose), modify, castStorablePtr,
   Struct,
   Record, Element, element,
   loadRecord, storeRecord, decomposeRecord, composeRecord,
   loadNewtype, storeNewtype, decomposeNewtype, composeNewtype,
   FirstClass, Stored,
   ) where

import LLVM.Extra.Class (MakeValueTuple, ValueTuple, Undefined, )
import qualified LLVM.Extra.ArithmeticPrivate as A
import qualified LLVM.Extra.Vector as Vector
import qualified LLVM.Extra.Array as Array

import qualified LLVM.Core as LLVM
import LLVM.Core
   (getElementPtr0,
    extractvalue, insertvalue,
    Value, -- valueOf, Vector,
    IsType, IsSized,
    CodeGenFunction, )
import LLVM.Util.Loop (Phi, )

import qualified Types.Data.Num as TypeNum
import Types.Data.Num (d0, d1, d2, )

import Foreign.StablePtr (StablePtr, )
import Foreign.Ptr (Ptr, castPtr, )

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

import Control.Monad (ap, )
import Control.Applicative (pure, liftA2, liftA3, )
import qualified Control.Applicative as App

import Data.Tuple.HT (fst3, snd3, thd3, )


{- |
An implementation of both 'MakeValueTuple' and 'Memory.C'
must ensure that @haskellValue@ is compatible
with @Stored (Struct haskellValue)@ (which we want to call @llvmStruct@).
That is, writing and reading @llvmStruct@ by LLVM
must be the same as accessing @haskellValue@ by 'Storable' methods.
ToDo: In future we may also require Storable constraint for @llvmStruct@.

We use a functional dependency in order to let type inference work nicely.
-}
class (Phi llvmValue, Undefined llvmValue, IsType (Struct llvmValue), IsSized (Struct llvmValue)) =>
      C llvmValue where
   type Struct llvmValue :: *
   load :: Value (Ptr (Struct llvmValue)) -> CodeGenFunction r llvmValue
   load ptr  =  decompose =<< LLVM.load ptr
   store :: llvmValue -> Value (Ptr (Struct llvmValue)) -> CodeGenFunction r ()
   store r ptr  =  flip LLVM.store ptr =<< compose r
   decompose :: Value (Struct llvmValue) -> CodeGenFunction r llvmValue
   compose :: llvmValue -> CodeGenFunction r (Value (Struct llvmValue))

modify ::
   (C llvmValue) =>
   (llvmValue -> CodeGenFunction r llvmValue) ->
   Value (Ptr (Struct llvmValue)) -> CodeGenFunction r ()
modify f ptr =
   flip store ptr =<< f =<< load ptr


type Record r o v = Element r o v v

data Element r o v x =
   Element {
      loadElement :: Value (Ptr o) -> CodeGenFunction r x,
      storeElement :: Value (Ptr o) -> v -> CodeGenFunction r (),
      extractElement :: Value o -> CodeGenFunction r x,
      insertElement :: v -> Value o -> CodeGenFunction r (Value o)
         -- State.Monoid
   }

element ::
   (C x,
    LLVM.GetValue o n, LLVM.ValueType o n ~ Struct x,
    LLVM.GetElementPtr o (n, ()), LLVM.ElementPtrType o (n, ()) ~ Struct x) =>
   (v -> x) -> n -> Element r o v x
element field n =
   Element {
      loadElement = \ptr -> load =<< getElementPtr0 ptr (n, ()),
      storeElement = \ptr v -> store (field v) =<< getElementPtr0 ptr (n, ()),
      extractElement = \o -> decompose =<< extractvalue o n,
      insertElement = \v o -> flip (insertvalue o) n =<< compose (field v)
   }

instance Functor (Element r o v) where
   fmap f m =
      Element {
         loadElement = fmap f . loadElement m,
         storeElement = storeElement m,
         extractElement = fmap f . extractElement m,
         insertElement = insertElement m
      }

instance App.Applicative (Element r o v) where
   pure x =
      Element {
         loadElement = \ _ptr -> return x,
         storeElement = \ _ptr _v -> return (),
         extractElement = \ _o -> return x,
         insertElement = \ _v o -> return o
      }
   f <*> x =
      Element {
         loadElement = \ptr -> loadElement f ptr `ap` loadElement x ptr,
         storeElement = \ptr y -> storeElement f ptr y >> storeElement x ptr y,
         extractElement = \o -> extractElement f o `ap` extractElement x o,
         insertElement = \y o -> insertElement f y o >>= insertElement x y
      }


loadRecord ::
   Record r o llvmValue ->
   Value (Ptr o) -> CodeGenFunction r llvmValue
loadRecord = loadElement

storeRecord ::
   Record r o llvmValue ->
   llvmValue -> Value (Ptr o) -> CodeGenFunction r ()
storeRecord m y ptr = storeElement m ptr y

decomposeRecord ::
   Record r o llvmValue ->
   Value o -> CodeGenFunction r llvmValue
decomposeRecord m =
   extractElement m

composeRecord ::
   (IsType o) =>
   Record r o llvmValue ->
   llvmValue -> CodeGenFunction r (Value o)
composeRecord m v =
   insertElement m v (LLVM.value LLVM.undef)



pair ::
   (C a, C b) =>
   Record r (LLVM.Struct (Struct a, (Struct b, ()))) (a, b)
pair =
   liftA2 (,)
      (element fst d0)
      (element snd d1)

instance (C a, C b) => C (a, b) where
   type Struct (a, b) = LLVM.Struct (Struct a, (Struct b, ()))
   load = loadRecord pair
   store = storeRecord pair
   decompose = decomposeRecord pair
   compose = composeRecord pair


triple ::
   (C a, C b, C c) =>
   Record r (LLVM.Struct (Struct a, (Struct b, (Struct c, ())))) (a, b, c)
triple =
   liftA3 (,,)
      (element fst3 d0)
      (element snd3 d1)
      (element thd3 d2)

instance (C a, C b, C c) => C (a, b, c) where
   type Struct (a, b, c) =
           LLVM.Struct (Struct a, (Struct b, (Struct c, ())))
   load = loadRecord triple
   store = storeRecord triple
   decompose = decomposeRecord triple
   compose = composeRecord triple


{-
This would not work for Booleans,
since on x86 LLVM's @i1@ type uses one byte in memory,
whereas Storable uses 4 byte and 4 byte alignment.

instance (LLVM.IsFirstClass a) => C (Value a) a where
   load = LLVM.load
   store = LLVM.store
   decompose = return
   compose = return
-}


class (LLVM.IsFirstClass llvmType, IsType (Stored llvmType)) =>
      FirstClass llvmType where
   type Stored llvmType :: *
   fromStorable :: Value (Stored llvmType) -> CodeGenFunction r (Value llvmType)
   toStorable :: Value llvmType -> CodeGenFunction r (Value (Stored llvmType))

instance FirstClass Float  where type Stored Float  = Float  ; fromStorable = return; toStorable = return
instance FirstClass Double where type Stored Double = Double ; fromStorable = return; toStorable = return
instance FirstClass Int8   where type Stored Int8   = Int8   ; fromStorable = return; toStorable = return
instance FirstClass Int16  where type Stored Int16  = Int16  ; fromStorable = return; toStorable = return
instance FirstClass Int32  where type Stored Int32  = Int32  ; fromStorable = return; toStorable = return
instance FirstClass Int64  where type Stored Int64  = Int64  ; fromStorable = return; toStorable = return
instance FirstClass Word8  where type Stored Word8  = Word8  ; fromStorable = return; toStorable = return
instance FirstClass Word16 where type Stored Word16 = Word16 ; fromStorable = return; toStorable = return
instance FirstClass Word32 where type Stored Word32 = Word32 ; fromStorable = return; toStorable = return
instance FirstClass Word64 where type Stored Word64 = Word64 ; fromStorable = return; toStorable = return
instance FirstClass Bool   where
   type Stored Bool = Word32
   fromStorable = A.cmp LLVM.CmpNE (LLVM.value LLVM.zero)
   toStorable = LLVM.zext
instance
   (TypeNum.PositiveT n, LLVM.IsPrimitive a, LLVM.IsPrimitive (Stored a), FirstClass a) =>
      FirstClass (LLVM.Vector n a) where
   type Stored (LLVM.Vector n a) = LLVM.Vector n (Stored a)
   fromStorable = Vector.map fromStorable
   toStorable = Vector.map toStorable
instance
   (TypeNum.NaturalT n, LLVM.IsFirstClass (Stored a),
    FirstClass a, IsSized a, IsSized (Stored a)) =>
      FirstClass (LLVM.Array n a) where
   type Stored (LLVM.Array n a) = LLVM.Array n (Stored a)
   fromStorable = Array.map fromStorable
   toStorable = Array.map toStorable

instance (IsType a) => FirstClass (Ptr a) where
   type Stored (Ptr a) = Ptr a
   fromStorable = return; toStorable = return
instance FirstClass (StablePtr a) where
   type Stored (StablePtr a) = StablePtr a
   fromStorable = return; toStorable = return


instance
   (LLVM.IsFirstClass (LLVM.Struct s),
    IsType (LLVM.Struct (StoredStruct s)),
    ConvertStruct s TypeNum.D0 s) =>
      FirstClass (LLVM.Struct s) where
   type Stored (LLVM.Struct s) = LLVM.Struct (StoredStruct s)
   fromStorable sm =
      case undefined of
         sfields -> do
            s <- decomposeField sfields d0 sm
            let _ = asTypeOf (fields s) sfields
            return s
   toStorable s =
      composeField (fields s) d0 s

fields :: Value (LLVM.Struct s) -> s
fields _ = undefined


type family StoredStruct s :: *
type instance StoredStruct () = ()
type instance StoredStruct (s,rem) = (Stored s, StoredStruct rem)

class
   ConvertStruct s i rem where
   decomposeField ::
      rem -> i -> Value (LLVM.Struct (StoredStruct s)) ->
      CodeGenFunction r (Value (LLVM.Struct s))
   composeField ::
      rem -> i -> Value (LLVM.Struct s) ->
      CodeGenFunction r (Value (LLVM.Struct (StoredStruct s)))

instance
   (sm ~ StoredStruct s,
    LLVM.GetValue (LLVM.Struct s) i, LLVM.ValueType (LLVM.Struct s) i ~ a,
    LLVM.GetValue (LLVM.Struct sm) i, LLVM.ValueType (LLVM.Struct sm) i ~ am,
    FirstClass a, am ~ Stored a,
    ConvertStruct s (TypeNum.Succ i) rem) =>
      ConvertStruct s i (a,rem) where
   decomposeField ~(_,rem_) i sm = do
      s <- decomposeField rem_ (TypeNum.succT i) sm
      a <- fromStorable =<< LLVM.extractvalue sm i
      LLVM.insertvalue s a i
   composeField ~(_,rem_) i s = do
      sm <- composeField rem_ (TypeNum.succT i) s
      am <- toStorable =<< LLVM.extractvalue s i
      LLVM.insertvalue sm am i

instance
   (sm ~ StoredStruct s,
    IsType (LLVM.Struct s),
    IsType (LLVM.Struct sm)) =>
      ConvertStruct s i () where
   decomposeField _ _ _ =
      return (LLVM.value LLVM.undef)
   composeField _ _ _ =
      return (LLVM.value LLVM.undef)


instance (FirstClass a, IsSized (Stored a)) => C (Value a) where
   type Struct (Value a) = Stored a
   decompose = fromStorable
   compose = toStorable


instance C () where
   type Struct () = LLVM.Struct ()
   load _ = return ()
   store _ _ = return ()
   decompose _ = return ()
   compose _ = return (LLVM.value LLVM.undef)

castStorablePtr ::
   (MakeValueTuple haskellValue, C (ValueTuple haskellValue)) =>
   Ptr haskellValue -> Ptr (Struct (ValueTuple haskellValue))
castStorablePtr = castPtr



loadNewtype ::
   (C a) =>
   (a -> llvmValue) ->
   Value (Ptr (Struct a)) -> CodeGenFunction r llvmValue
loadNewtype wrap ptr =
   fmap wrap $ load ptr

storeNewtype ::
   (C a) =>
   (llvmValue -> a) ->
   llvmValue -> Value (Ptr (Struct a)) -> CodeGenFunction r ()
storeNewtype unwrap y ptr =
   store (unwrap y) ptr

decomposeNewtype ::
   (C a) =>
   (a -> llvmValue) ->
   Value (Struct a) -> CodeGenFunction r llvmValue
decomposeNewtype wrap y =
   fmap wrap $ decompose y

composeNewtype ::
   (C a) =>
   (llvmValue -> a) ->
   llvmValue -> CodeGenFunction r (Value (Struct a))
composeNewtype unwrap y =
   compose (unwrap y)