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

import GHC.TypeLits
import Control.Monad.Primitive
import Data.Primitive.Types
import Data.Primitive.ByteArray

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
               , IsMultipleOf n (SizeOf a)
               ) => 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
               , IsMultipleOf n (SizeOf a)
               ) => MutableByteArray (PrimState m) -> Offset n -> m a
readAtByte buf (Offset offset) = readByteArray buf (div offset $ (sizeOf (undefined :: a)))
{-# INLINE readAtByte #-}