-- |
-- Module      : Data.ByteArray.Types
-- License     : BSD-style
-- Maintainer  : Vincent Hanquez <vincent@snarc.org>
-- Stability   : stable
-- Portability : Good
--
{-# LANGUAGE CPP #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE Rank2Types    #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE TypeFamilies  #-}
{-# LANGUAGE UndecidableInstances #-}
module Data.ByteArray.Types
    ( ByteArrayAccess(..)
    , ByteArray(..)
    ) where

import           Foreign.Ptr
import           Data.Monoid

#ifdef WITH_BYTESTRING_SUPPORT
import qualified Data.ByteString as Bytestring (length)
import qualified Data.ByteString.Internal as Bytestring
import           Foreign.ForeignPtr (withForeignPtr)
#endif

import           Data.Memory.PtrMethods (memCopy)


import           Data.Proxy (Proxy(..))
import           Data.Word (Word8)

import qualified Basement.Types.OffsetSize as Base
import qualified Basement.UArray as Base
import qualified Basement.String as Base (String, toBytes, Encoding(UTF8))

import qualified Basement.UArray.Mutable as BaseMutable (withMutablePtrHint)
import qualified Basement.Block as Block
import qualified Basement.Block.Mutable as Block

import           Basement.Nat
import qualified Basement.Sized.Block as BlockN

import Prelude hiding (length)

-- | Class to Access size properties and data of a ByteArray
class ByteArrayAccess ba where
    -- | Return the length in bytes of a bytearray
    length        :: ba -> Int
    -- | Allow to use using a pointer
    withByteArray :: ba -> (Ptr p -> IO a) -> IO a
    -- | Copy the data of a bytearray to a ptr
    copyByteArrayToPtr :: ba -> Ptr p -> IO ()
    copyByteArrayToPtr ba
a Ptr p
dst = ba -> (Ptr Word8 -> IO ()) -> IO ()
forall ba p a. ByteArrayAccess ba => ba -> (Ptr p -> IO a) -> IO a
withByteArray ba
a ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
src -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
memCopy (Ptr p -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr p
dst) Ptr Word8
src (ba -> Int
forall ba. ByteArrayAccess ba => ba -> Int
length ba
a)

-- | Class to allocate new ByteArray of specific size
class (Eq ba, Ord ba, Monoid ba, ByteArrayAccess ba) => ByteArray ba where
    -- | allocate `n` bytes and perform the given operation
    allocRet  :: Int
                -- ^ number of bytes to allocate. i.e. might not match the
                -- size of the given type `ba`.
              -> (Ptr p -> IO a)
              -> IO (a, ba)

#ifdef WITH_BYTESTRING_SUPPORT
instance ByteArrayAccess Bytestring.ByteString where
    length :: ByteString -> Int
length = ByteString -> Int
Bytestring.length
    withByteArray :: ByteString -> (Ptr p -> IO a) -> IO a
withByteArray (Bytestring.PS ForeignPtr Word8
fptr Int
off Int
_) Ptr p -> IO a
f = ForeignPtr Word8 -> (Ptr Word8 -> IO a) -> IO a
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fptr ((Ptr Word8 -> IO a) -> IO a) -> (Ptr Word8 -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> Ptr p -> IO a
f (Ptr p -> IO a) -> Ptr p -> IO a
forall a b. (a -> b) -> a -> b
$! (Ptr Word8
ptr Ptr Word8 -> Int -> Ptr p
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
off)

instance ByteArray Bytestring.ByteString where
    allocRet :: Int -> (Ptr p -> IO a) -> IO (a, ByteString)
allocRet Int
sz Ptr p -> IO a
f = do
        ForeignPtr Word8
fptr <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
Bytestring.mallocByteString Int
sz
        a
r    <- ForeignPtr Word8 -> (Ptr Word8 -> IO a) -> IO a
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fptr (Ptr p -> IO a
f (Ptr p -> IO a) -> (Ptr Word8 -> Ptr p) -> Ptr Word8 -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr Word8 -> Ptr p
forall a b. Ptr a -> Ptr b
castPtr)
        (a, ByteString) -> IO (a, ByteString)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
r, ForeignPtr Word8 -> Int -> Int -> ByteString
Bytestring.PS ForeignPtr Word8
fptr Int
0 Int
sz)
#endif

#ifdef WITH_BASEMENT_SUPPORT

baseBlockRecastW8 :: Base.PrimType ty => Block.Block ty -> Block.Block Word8
baseBlockRecastW8 :: Block ty -> Block Word8
baseBlockRecastW8 = Block ty -> Block Word8
forall b a. PrimType b => Block a -> Block b
Block.unsafeCast -- safe with Word8 destination

instance Base.PrimType ty => ByteArrayAccess (Block.Block ty) where
    length :: Block ty -> Int
length Block ty
a = let Base.CountOf Int
i = Block Word8 -> CountOf Word8
forall ty. PrimType ty => Block ty -> CountOf ty
Block.length (Block ty -> Block Word8
forall ty. PrimType ty => Block ty -> Block Word8
baseBlockRecastW8 Block ty
a) in Int
i
    withByteArray :: Block ty -> (Ptr p -> IO a) -> IO a
withByteArray Block ty
a Ptr p -> IO a
f = Block Word8 -> (Ptr Word8 -> IO a) -> IO a
forall (prim :: * -> *) ty a.
PrimMonad prim =>
Block ty -> (Ptr ty -> prim a) -> prim a
Block.withPtr (Block ty -> Block Word8
forall ty. PrimType ty => Block ty -> Block Word8
baseBlockRecastW8 Block ty
a) (Ptr p -> IO a
f (Ptr p -> IO a) -> (Ptr Word8 -> Ptr p) -> Ptr Word8 -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr Word8 -> Ptr p
forall a b. Ptr a -> Ptr b
castPtr)
    copyByteArrayToPtr :: Block ty -> Ptr p -> IO ()
copyByteArrayToPtr Block ty
ba Ptr p
dst = do
        MutableBlock Word8 RealWorld
mb <- Block Word8 -> IO (MutableBlock Word8 (PrimState IO))
forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
Block ty -> prim (MutableBlock ty (PrimState prim))
Block.unsafeThaw (Block ty -> Block Word8
forall ty. PrimType ty => Block ty -> Block Word8
baseBlockRecastW8 Block ty
ba)
        MutableBlock Word8 (PrimState IO)
-> Offset Word8 -> Ptr Word8 -> CountOf Word8 -> IO ()
forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
MutableBlock ty (PrimState prim)
-> Offset ty -> Ptr ty -> CountOf ty -> prim ()
Block.copyToPtr MutableBlock Word8 RealWorld
MutableBlock Word8 (PrimState IO)
mb Offset Word8
0 (Ptr p -> Ptr Word8
forall a b. Ptr a -> Ptr b
castPtr Ptr p
dst) (Block Word8 -> CountOf Word8
forall ty. PrimType ty => Block ty -> CountOf ty
Block.length (Block Word8 -> CountOf Word8) -> Block Word8 -> CountOf Word8
forall a b. (a -> b) -> a -> b
$ Block ty -> Block Word8
forall ty. PrimType ty => Block ty -> Block Word8
baseBlockRecastW8 Block ty
ba)

instance (KnownNat n, Base.PrimType ty, Base.Countable ty n) => ByteArrayAccess (BlockN.BlockN n ty) where
    length :: BlockN n ty -> Int
length BlockN n ty
a = let Base.CountOf Int
i = BlockN n ty -> CountOf Word8
forall (n :: Nat) ty. PrimType ty => BlockN n ty -> CountOf Word8
BlockN.lengthBytes BlockN n ty
a in Int
i
    withByteArray :: BlockN n ty -> (Ptr p -> IO a) -> IO a
withByteArray BlockN n ty
a Ptr p -> IO a
f = BlockN n ty -> (Ptr ty -> IO a) -> IO a
forall (prim :: * -> *) (n :: Nat) ty a.
(PrimMonad prim, KnownNat n) =>
BlockN n ty -> (Ptr ty -> prim a) -> prim a
BlockN.withPtr BlockN n ty
a (Ptr p -> IO a
f (Ptr p -> IO a) -> (Ptr ty -> Ptr p) -> Ptr ty -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr ty -> Ptr p
forall a b. Ptr a -> Ptr b
castPtr)
    copyByteArrayToPtr :: BlockN n ty -> Ptr p -> IO ()
copyByteArrayToPtr BlockN n ty
bna = Block ty -> Ptr p -> IO ()
forall ba p. ByteArrayAccess ba => ba -> Ptr p -> IO ()
copyByteArrayToPtr (BlockN n ty -> Block ty
forall (n :: Nat) ty. BlockN n ty -> Block ty
BlockN.toBlock BlockN n ty
bna)

baseUarrayRecastW8 :: Base.PrimType ty => Base.UArray ty -> Base.UArray Word8
baseUarrayRecastW8 :: UArray ty -> UArray Word8
baseUarrayRecastW8 = UArray ty -> UArray Word8
forall a b. (PrimType a, PrimType b) => UArray a -> UArray b
Base.recast

instance Base.PrimType ty => ByteArrayAccess (Base.UArray ty) where
    length :: UArray ty -> Int
length UArray ty
a = let Base.CountOf Int
i = UArray Word8 -> CountOf Word8
forall ty. UArray ty -> CountOf ty
Base.length (UArray ty -> UArray Word8
forall ty. PrimType ty => UArray ty -> UArray Word8
baseUarrayRecastW8 UArray ty
a) in Int
i
    withByteArray :: UArray ty -> (Ptr p -> IO a) -> IO a
withByteArray UArray ty
a Ptr p -> IO a
f = UArray Word8 -> (Ptr Word8 -> IO a) -> IO a
forall ty (prim :: * -> *) a.
(PrimMonad prim, PrimType ty) =>
UArray ty -> (Ptr ty -> prim a) -> prim a
Base.withPtr (UArray ty -> UArray Word8
forall ty. PrimType ty => UArray ty -> UArray Word8
baseUarrayRecastW8 UArray ty
a) (Ptr p -> IO a
f (Ptr p -> IO a) -> (Ptr Word8 -> Ptr p) -> Ptr Word8 -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr Word8 -> Ptr p
forall a b. Ptr a -> Ptr b
castPtr)
    copyByteArrayToPtr :: UArray ty -> Ptr p -> IO ()
copyByteArrayToPtr UArray ty
ba Ptr p
dst = UArray ty -> Ptr ty -> IO ()
forall ty (prim :: * -> *).
(PrimType ty, PrimMonad prim) =>
UArray ty -> Ptr ty -> prim ()
Base.copyToPtr UArray ty
ba (Ptr p -> Ptr ty
forall a b. Ptr a -> Ptr b
castPtr Ptr p
dst)

instance ByteArrayAccess Base.String where
    length :: String -> Int
length String
str = let Base.CountOf Int
i = UArray Word8 -> CountOf Word8
forall ty. UArray ty -> CountOf ty
Base.length UArray Word8
bytes in Int
i
      where
        -- the Foundation's length return a number of elements not a number of
        -- bytes. For @ByteArrayAccess@, because we are using an @Int@, we
        -- didn't see that we were returning the wrong @CountOf@.
        bytes :: UArray Word8
bytes = Encoding -> String -> UArray Word8
Base.toBytes Encoding
Base.UTF8 String
str
    withByteArray :: String -> (Ptr p -> IO a) -> IO a
withByteArray String
s Ptr p -> IO a
f = UArray Word8 -> (Ptr p -> IO a) -> IO a
forall ba p a. ByteArrayAccess ba => ba -> (Ptr p -> IO a) -> IO a
withByteArray (Encoding -> String -> UArray Word8
Base.toBytes Encoding
Base.UTF8 String
s) Ptr p -> IO a
f

instance (Ord ty, Base.PrimType ty) => ByteArray (Block.Block ty) where
    allocRet :: Int -> (Ptr p -> IO a) -> IO (a, Block ty)
allocRet Int
sz Ptr p -> IO a
f = do
        MutableBlock ty RealWorld
mba <- CountOf ty -> IO (MutableBlock ty (PrimState IO))
forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
CountOf ty -> prim (MutableBlock ty (PrimState prim))
Block.new (CountOf ty -> IO (MutableBlock ty (PrimState IO)))
-> CountOf ty -> IO (MutableBlock ty (PrimState IO))
forall a b. (a -> b) -> a -> b
$ Int -> Proxy ty -> CountOf ty
forall ty. PrimType ty => Int -> Proxy ty -> CountOf ty
sizeRecastBytes Int
sz Proxy ty
forall k (t :: k). Proxy t
Proxy
        a
a   <- Bool
-> Bool
-> MutableBlock ty (PrimState IO)
-> (Ptr ty -> IO a)
-> IO a
forall ty (prim :: * -> *) a.
PrimMonad prim =>
Bool
-> Bool
-> MutableBlock ty (PrimState prim)
-> (Ptr ty -> prim a)
-> prim a
Block.withMutablePtrHint Bool
True Bool
False MutableBlock ty RealWorld
MutableBlock ty (PrimState IO)
mba (Ptr p -> IO a
f (Ptr p -> IO a) -> (Ptr ty -> Ptr p) -> Ptr ty -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr ty -> Ptr p
forall a b. Ptr a -> Ptr b
castPtr)
        Block ty
ba  <- MutableBlock ty (PrimState IO) -> IO (Block ty)
forall (prim :: * -> *) ty.
PrimMonad prim =>
MutableBlock ty (PrimState prim) -> prim (Block ty)
Block.unsafeFreeze MutableBlock ty RealWorld
MutableBlock ty (PrimState IO)
mba
        (a, Block ty) -> IO (a, Block ty)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, Block ty
ba)

instance (Ord ty, Base.PrimType ty) => ByteArray (Base.UArray ty) where
    allocRet :: Int -> (Ptr p -> IO a) -> IO (a, UArray ty)
allocRet Int
sz Ptr p -> IO a
f = do
        MUArray ty RealWorld
mba <- CountOf ty -> IO (MUArray ty (PrimState IO))
forall (prim :: * -> *) ty.
(PrimMonad prim, PrimType ty) =>
CountOf ty -> prim (MUArray ty (PrimState prim))
Base.new (CountOf ty -> IO (MUArray ty (PrimState IO)))
-> CountOf ty -> IO (MUArray ty (PrimState IO))
forall a b. (a -> b) -> a -> b
$ Int -> Proxy ty -> CountOf ty
forall ty. PrimType ty => Int -> Proxy ty -> CountOf ty
sizeRecastBytes Int
sz Proxy ty
forall k (t :: k). Proxy t
Proxy
        a
a   <- Bool
-> Bool -> MUArray ty (PrimState IO) -> (Ptr ty -> IO a) -> IO a
forall ty (prim :: * -> *) a.
(PrimMonad prim, PrimType ty) =>
Bool
-> Bool
-> MUArray ty (PrimState prim)
-> (Ptr ty -> prim a)
-> prim a
BaseMutable.withMutablePtrHint Bool
True Bool
False MUArray ty RealWorld
MUArray ty (PrimState IO)
mba (Ptr p -> IO a
f (Ptr p -> IO a) -> (Ptr ty -> Ptr p) -> Ptr ty -> IO a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Ptr ty -> Ptr p
forall a b. Ptr a -> Ptr b
castPtr)
        UArray ty
ba  <- MUArray ty (PrimState IO) -> IO (UArray ty)
forall (prim :: * -> *) ty.
PrimMonad prim =>
MUArray ty (PrimState prim) -> prim (UArray ty)
Base.unsafeFreeze MUArray ty RealWorld
MUArray ty (PrimState IO)
mba
        (a, UArray ty) -> IO (a, UArray ty)
forall (m :: * -> *) a. Monad m => a -> m a
return (a
a, UArray ty
ba)

sizeRecastBytes :: Base.PrimType ty => Int -> Proxy ty -> Base.CountOf ty
sizeRecastBytes :: Int -> Proxy ty -> CountOf ty
sizeRecastBytes Int
w Proxy ty
p = Int -> CountOf ty
forall ty. Int -> CountOf ty
Base.CountOf (Int -> CountOf ty) -> Int -> CountOf ty
forall a b. (a -> b) -> a -> b
$
    let (Int
q,Int
r) = Int
w Int -> Int -> (Int, Int)
forall a. Integral a => a -> a -> (a, a)
`Prelude.quotRem` Int
szTy
     in Int
q Int -> Int -> Int
forall a. Num a => a -> a -> a
+ (if Int
r Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0 then Int
0 else Int
1)
  where !(Base.CountOf Int
szTy) = Proxy ty -> CountOf Word8
forall ty. PrimType ty => Proxy ty -> CountOf Word8
Base.primSizeInBytes Proxy ty
p
{-# INLINE [1] sizeRecastBytes #-}

#endif