{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE KindSignatures        #-}
{-# LANGUAGE MagicHash             #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE ScopedTypeVariables   #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE CPP                   #-}
module Data.Primitive.ByteArray.Util where

import GHC.Exts (Ptr(..))
import GHC.TypeLits
import Control.Monad.Primitive
import qualified Data.ByteString as B
import qualified Data.ByteString.Unsafe as B
import Data.Primitive.Addr
import Data.Primitive.Types
import Data.Primitive.ByteArray
import Data.Word
import Foreign.Ptr

newtype Offset (n :: Nat) = Offset Int

offset :: KnownNat n => Offset n
offset = let k = Offset $ fromIntegral $ natVal k in k

(+.) :: Offset j -> Offset k -> Offset (j + k)
Offset j +. Offset k = Offset (j + k)


type family SizeOf a :: Nat

class IsMultipleOf (n::Nat) (k::Nat)

instance n ~ (q * k) => IsMultipleOf n k

writeAtByte :: ( Prim a
               , PrimMonad m
#if __GLASGOW_HASKELL__ >= 802
               , IsMultipleOf n (SizeOf a)
#endif
               ) => MutableByteArray (PrimState m) -> Offset n -> a -> m ()
writeAtByte buf (Offset offset) a = writeByteArray buf (div offset $ (sizeOf a)) a
{-# INLINE writeAtByte #-}

readAtByte :: forall a m n.
               ( Prim a
               , PrimMonad m
#if __GLASGOW_HASKELL__ >= 802
               , IsMultipleOf n (SizeOf a)
#endif
               ) => MutableByteArray (PrimState m) -> Offset n -> m a
readAtByte buf (Offset offset) = readByteArray buf (div offset $ (sizeOf (undefined :: a)))
{-# INLINE readAtByte #-}

writeStringAt :: PrimMonad m => MutableByteArray s -> Int -> B.ByteString -> m (Ptr a)
writeStringAt src o bsname = do
    (p,cnt) <- unsafeIOToPrim $ B.unsafeUseAsCStringLen bsname return
    let nptr = ptr (mutableByteArrayContents src) `plusPtr` o
    copyAddr (adr nptr) (adr p) cnt
    writeOffAddr (adr nptr) cnt (0 :: Word8)
    return nptr

ptr :: Addr -> Ptr a
ptr (Addr a) = Ptr a

adr :: Ptr a -> Addr
adr (Ptr a) = Addr a