summaryrefslogtreecommitdiff
path: root/src/Data/Primitive/ByteArray/Util.hs
blob: 83a2e7d2e5cab7c4e7bdfb3a420a113c0ae0aed6 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
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