1 files changed, 71 insertions, 0 deletions

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