1 files changed, 45 insertions, 0 deletions

diff --git a/haskell/Data/Primitive/ByteArray/Util.hs b/haskell/Data/Primitive/ByteArray/Util.hs
new file mode 100644
index 0000000..1776286
--- /dev/null
+++ b/haskell/Data/Primitive/ByteArray/Util.hs
@@ -0,0 +1,45 @@
+{-# 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 #-}