{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE MagicHash #-}
-- |
-- Module      : Data.Prim.Memory.ByteString
-- Copyright   : (c) Alexey Kuleshevich 2020
-- License     : BSD3
-- Maintainer  : Alexey Kuleshevich <alexey@kuleshevi.ch>
-- Stability   : experimental
-- Portability : non-portable
--
module Data.Prim.Memory.ByteString
  (
    MByteString(..)
  -- * Conversion
  -- Builder
  , Builder
  , toBuilderBytes
  , fromBuilderBytes
  -- ** ByteString
  , ByteString(..)
  , toByteStringBytes
  , fromByteStringBytes
  , fromLazyByteStringBytes
  , withPtrByteString
  , withNoHaltPtrByteString
  -- ** ShortByteString
  , ShortByteString(..)
  , toShortByteStringBytes
  , fromShortByteStringBytes
  , byteStringConvertError
  ) where

import Control.Monad.ST
import Data.ByteString.Builder
import Data.ByteString.Internal
import Data.ByteString.Short.Internal
import qualified Data.ByteString as BS
import qualified Data.ByteString.Lazy as BSL
import Data.Prim
import Foreign.Prim
import Control.Prim.Monad
import Control.Prim.Eval
import GHC.ForeignPtr
import Data.Prim.Memory.Ptr
import Data.Prim.Memory.Bytes.Internal
  ( Bytes(..)
  , Pinned(..)
  , allocMBytes
  , freezeMBytes
  , byteCountBytes
  , relaxPinnedBytes
  , toForeignPtrBytes
  , castForeignPtrToBytes
  , byteStringConvertError
  )

-- | Mutable version of a `ByteString`
newtype MByteString s = MByteString ByteString


-- | /O(1)/ - Cast immutable `Bytes` to an immutable `ByteString`
--
-- @since 0.1.0
toByteStringBytes :: Bytes 'Pin -> ByteString
{-# INLINE toByteStringBytes #-}
toByteStringBytes :: Bytes 'Pin -> ByteString
toByteStringBytes Bytes 'Pin
b =
#if MIN_VERSION_bytestring(0,11,0)
  BS (toForeignPtrBytes b) (coerce (byteCountBytes b))
#else
  ForeignPtr Word8 -> Int -> Int -> ByteString
PS (Bytes 'Pin -> ForeignPtr Word8
forall e. Bytes 'Pin -> ForeignPtr e
toForeignPtrBytes Bytes 'Pin
b) Int
0 (Count Word8 -> Int
coerce (Bytes 'Pin -> Count Word8
forall (p :: Pinned). Bytes p -> Count Word8
byteCountBytes Bytes 'Pin
b))
#endif


-- | /O(1)/ - Cast an immutable `ByteString` to immutable `Bytes`. Only unsliced
-- `ByteString`s that are backed by a `ForeignPtr` allocated on Haskell heap without
-- finilizers can be converted without copy.
--
-- @since 0.2.0
castByteStringBytes :: ByteString -> Either String (Bytes 'Pin)
#if MIN_VERSION_bytestring(0,11,0)
castByteStringBytes (BS fptr n) = do
#else
castByteStringBytes :: ByteString -> Either String (Bytes 'Pin)
castByteStringBytes (PS ForeignPtr Word8
fptr Int
o Int
n) = do
  Bool -> Either String () -> Either String ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Int
o Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
0) Either String ()
forall b. Either String b
sliceError
#endif
  Bytes 'Pin
b <- ForeignPtr Word8 -> Either String (Bytes 'Pin)
forall e. ForeignPtr e -> Either String (Bytes 'Pin)
castForeignPtrToBytes ForeignPtr Word8
fptr
  Bool -> Either String () -> Either String ()
forall (f :: * -> *). Applicative f => Bool -> f () -> f ()
unless (Count Word8 -> Int
forall e. Count e -> Int
unCount (Bytes 'Pin -> Count Word8
forall (p :: Pinned). Bytes p -> Count Word8
byteCountBytes Bytes 'Pin
b) Int -> Int -> Bool
forall a. Eq a => a -> a -> Bool
== Int
n) Either String ()
forall b. Either String b
sliceError
  Bytes 'Pin -> Either String (Bytes 'Pin)
forall a b. b -> Either a b
Right Bytes 'Pin
b
  where
    sliceError :: Either String b
sliceError = String -> Either String b
forall a b. a -> Either a b
Left String
"ByteString was sliced"
{-# INLINE castByteStringBytes #-}


-- | /O(1)/ - Cast an immutable `Bytes` to an immutable `ShortByteString`
--
-- @since 0.1.0
toShortByteStringBytes :: Bytes p -> ShortByteString
toShortByteStringBytes :: Bytes p -> ShortByteString
toShortByteStringBytes (Bytes ByteArray#
ba#) = ByteArray# -> ShortByteString
SBS ByteArray#
ba#
{-# INLINE toShortByteStringBytes #-}

-- | /O(1)/ - Cast an immutable  `ShortByteString` to an immutable `Bytes`
--
-- @since 0.1.0
fromShortByteStringBytes :: ShortByteString -> Bytes 'Inc
fromShortByteStringBytes :: ShortByteString -> Bytes 'Inc
fromShortByteStringBytes (SBS ByteArray#
ba#) = ByteArray# -> Bytes 'Inc
forall (p :: Pinned). ByteArray# -> Bytes p
Bytes ByteArray#
ba#
{-# INLINE fromShortByteStringBytes #-}

-- | Convert `Bytes` into a bytestring `Builder`
toBuilderBytes :: Bytes p -> Builder
toBuilderBytes :: Bytes p -> Builder
toBuilderBytes = ShortByteString -> Builder
shortByteString (ShortByteString -> Builder)
-> (Bytes p -> ShortByteString) -> Bytes p -> Builder
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Bytes p -> ShortByteString
forall (p :: Pinned). Bytes p -> ShortByteString
toShortByteStringBytes
{-# INLINE[1] toBuilderBytes #-}
{-# RULES
"toBuilderBytes" toBuilderBytes = byteString . toByteStringBytes
  #-}

-- | /O(n)/ - Allocate `Bytes` and fill them using the supplied `Builder`
fromBuilderBytes :: Builder -> Bytes 'Pin
fromBuilderBytes :: Builder -> Bytes 'Pin
fromBuilderBytes Builder
b = ByteString -> Bytes 'Pin
fromLazyByteStringBytes (Builder -> ByteString
toLazyByteString Builder
b)
{-# INLINE fromBuilderBytes #-}


-- | /O(n)/ - Allocate `Bytes` and fill them with the contents of a lazy `BSL.ByteString`
fromLazyByteStringBytes :: BSL.ByteString -> Bytes 'Pin
fromLazyByteStringBytes :: ByteString -> Bytes 'Pin
fromLazyByteStringBytes = ByteString -> Bytes 'Pin
forall (p :: Pinned). Typeable p => ByteString -> Bytes p
fromByteStringBytes (ByteString -> Bytes 'Pin)
-> (ByteString -> ByteString) -> ByteString -> Bytes 'Pin
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ByteString -> ByteString
BSL.toStrict
{-# INLINE fromLazyByteStringBytes #-}


-- | /O(n)/ - Convert a strict `ByteString` to `Bytes`.
fromByteStringBytes :: Typeable p => ByteString -> Bytes p
fromByteStringBytes :: ByteString -> Bytes p
fromByteStringBytes ByteString
bs =
  case ByteString -> Either String (Bytes 'Pin)
castByteStringBytes ByteString
bs of
    Right Bytes 'Pin
b -> Bytes 'Pin -> Bytes p
forall (p :: Pinned). Bytes 'Pin -> Bytes p
relaxPinnedBytes Bytes 'Pin
b
    Left String
_ ->
      (forall s. ST s (Bytes p)) -> Bytes p
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Bytes p)) -> Bytes p)
-> (forall s. ST s (Bytes p)) -> Bytes p
forall a b. (a -> b) -> a -> b
$
      ByteString -> (Ptr Word8 -> ST s (Bytes p)) -> ST s (Bytes p)
forall s (m :: * -> *) a b.
MonadPrim s m =>
ByteString -> (Ptr a -> m b) -> m b
withPtrByteString ByteString
bs ((Ptr Word8 -> ST s (Bytes p)) -> ST s (Bytes p))
-> (Ptr Word8 -> ST s (Bytes p)) -> ST s (Bytes p)
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
ptr -> do
        let c :: Count Word8
c = Int -> Count Word8
forall e. Int -> Count e
Count (ByteString -> Int
BS.length ByteString
bs) :: Count Word8
        MBytes p s
mb <- Count Word8 -> ST s (MBytes p s)
forall (p :: Pinned) e s (m :: * -> *).
(Typeable p, Prim e, MonadPrim s m) =>
Count e -> m (MBytes p s)
allocMBytes Count Word8
c
        Ptr Word8
-> Off Word8 -> MBytes p s -> Off Word8 -> Count Word8 -> ST s ()
forall s (m :: * -> *) e (p :: Pinned).
(MonadPrim s m, Prim e) =>
Ptr e -> Off e -> MBytes p s -> Off e -> Count e -> m ()
copyPtrToMBytes Ptr Word8
ptr Off Word8
0 MBytes p s
mb Off Word8
0 Count Word8
c
        MBytes p s -> ST s (Bytes p)
forall s (m :: * -> *) (p :: Pinned).
MonadPrim s m =>
MBytes p s -> m (Bytes p)
freezeMBytes MBytes p s
mb
{-# INLINE fromByteStringBytes #-}


withPtrByteString :: MonadPrim s m => ByteString -> (Ptr a -> m b) -> m b
#if MIN_VERSION_bytestring(0,11,0)
withPtrByteString (BS (ForeignPtr addr# ptrContents) _) f = do
#else
withPtrByteString :: ByteString -> (Ptr a -> m b) -> m b
withPtrByteString (PS (ForeignPtr Addr#
addr'# ForeignPtrContents
ptrContents) (I# Int#
o#) Int
_) Ptr a -> m b
f = do
  let addr# :: Addr#
addr# = Addr#
addr'# Addr# -> Int# -> Addr#
`plusAddr#` Int#
o#
#endif
  b
r <- Ptr a -> m b
f (Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
addr#)
  b
r b -> m () -> m b
forall (f :: * -> *) a b. Functor f => a -> f b -> f a
<$ ForeignPtrContents -> m ()
forall s (m :: * -> *) a. MonadPrim s m => a -> m ()
touch ForeignPtrContents
ptrContents
{-# INLINE withPtrByteString #-}


withNoHaltPtrByteString :: MonadUnliftPrim s m => ByteString -> (Ptr a -> m b) -> m b
#if MIN_VERSION_bytestring(0,11,0)
withNoHaltPtrByteString (BS (ForeignPtr addr# ptrContents) _) f = do
#else
withNoHaltPtrByteString :: ByteString -> (Ptr a -> m b) -> m b
withNoHaltPtrByteString (PS (ForeignPtr Addr#
addr'# ForeignPtrContents
ptrContents) (I# Int#
o#) Int
_) Ptr a -> m b
f = do
  let addr# :: Addr#
addr# = Addr#
addr'# Addr# -> Int# -> Addr#
`plusAddr#` Int#
o#
#endif
  ForeignPtrContents -> m b -> m b
forall s (m :: * -> *) a b. MonadUnliftPrim s m => a -> m b -> m b
keepAlive ForeignPtrContents
ptrContents (m b -> m b) -> m b -> m b
forall a b. (a -> b) -> a -> b
$ Ptr a -> m b
f (Addr# -> Ptr a
forall a. Addr# -> Ptr a
Ptr Addr#
addr#)
{-# INLINE withNoHaltPtrByteString #-}