{-# LANGUAGE UnboxedTuples #-}

-- | raehik's bytestring extras (reimplementations of unexported internals).

module Tmp.BSExt
  ( module Tmp.BSExt
  , B.mkDeferredByteString
  ) where

import GHC.ForeignPtr ( ForeignPtr, unsafeWithForeignPtr, withForeignPtr )
import Foreign.Ptr ( Ptr )
import Foreign.Marshal.Utils ( copyBytes )
import Data.ByteString.Internal qualified as B
import Data.ByteString ( ByteString )
import Data.Word ( Word8 )
import Control.Exception ( assert )
import GHC.IO ( IO(IO) )
import GHC.Exts ( runRW# )

-- | Copy the given number of bytes from the second area (source) into the first
--   (destination); the copied areas may not overlap.
--
-- Reimplemented from the unexported function
-- 'Data.ByteString.Internal.Type.memcpyFp'.
memcpyFp :: ForeignPtr Word8 -> ForeignPtr Word8 -> Int -> IO ()
memcpyFp :: ForeignPtr Word8 -> ForeignPtr Word8 -> Int -> IO ()
memcpyFp ForeignPtr Word8
fp ForeignPtr Word8
fq Int
s = ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr ForeignPtr Word8
fp ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
p ->
                     ForeignPtr Word8 -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
unsafeWithForeignPtr ForeignPtr Word8
fq ((Ptr Word8 -> IO ()) -> IO ()) -> (Ptr Word8 -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
q -> Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
p Ptr Word8
q Int
s

-- | Create a 'ByteString' of size @l@ and use action @f@ to fill its contents.
--
-- Reimplemented from the unexported function
-- 'Data.ByteString.Internal.Type.createFp.
createFp :: Int -> (ForeignPtr Word8 -> IO ()) -> IO ByteString
createFp :: Int -> (ForeignPtr Word8 -> IO ()) -> IO ByteString
createFp Int
len ForeignPtr Word8 -> IO ()
action = Bool -> IO ByteString -> IO ByteString
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO ByteString -> IO ByteString) -> IO ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ do
    ForeignPtr Word8
fp <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
B.mallocByteString Int
len
    ForeignPtr Word8 -> IO ()
action ForeignPtr Word8
fp
    ForeignPtr Word8 -> Int -> IO ByteString
B.mkDeferredByteString ForeignPtr Word8
fp Int
len
{-# INLINE createFp #-}

createUptoNCPS
    :: Int
    -> (Ptr Word8 -> (Int -> IO ByteString) -> IO r)
    -> IO r
createUptoNCPS :: forall r.
Int -> (Ptr Word8 -> (Int -> IO ByteString) -> IO r) -> IO r
createUptoNCPS Int
maxLen Ptr Word8 -> (Int -> IO ByteString) -> IO r
action = Bool -> IO r -> IO r
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
maxLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO r -> IO r) -> IO r -> IO r
forall a b. (a -> b) -> a -> b
$ do
    ForeignPtr Word8
fp <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
B.mallocByteString Int
maxLen
    ForeignPtr Word8 -> (Ptr Word8 -> IO r) -> IO r
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fp ((Ptr Word8 -> IO r) -> IO r) -> (Ptr Word8 -> IO r) -> IO r
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
p -> Ptr Word8 -> (Int -> IO ByteString) -> IO r
action Ptr Word8
p ((Int -> IO ByteString) -> IO r) -> (Int -> IO ByteString) -> IO r
forall a b. (a -> b) -> a -> b
$ \Int
len ->
        ForeignPtr Word8 -> Int -> IO ByteString
B.mkDeferredByteString ForeignPtr Word8
fp Int
len
{-# INLINE createUptoNCPS #-}

createCPS
    :: Int
    -> (ForeignPtr Word8 -> Int -> IO r)
    -> ((Int -> IO r) -> Ptr Word8 -> IO r)
    -> IO r
createCPS :: forall r.
Int
-> (ForeignPtr Word8 -> Int -> IO r)
-> ((Int -> IO r) -> Ptr Word8 -> IO r)
-> IO r
createCPS Int
maxLen ForeignPtr Word8 -> Int -> IO r
finalize (Int -> IO r) -> Ptr Word8 -> IO r
f = Bool -> IO r -> IO r
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
maxLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO r -> IO r) -> IO r -> IO r
forall a b. (a -> b) -> a -> b
$ do
    ForeignPtr Word8
fp <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
B.mallocByteString Int
maxLen
    ForeignPtr Word8 -> (Ptr Word8 -> IO r) -> IO r
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fp ((Ptr Word8 -> IO r) -> IO r) -> (Ptr Word8 -> IO r) -> IO r
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
buf -> (Int -> IO r) -> Ptr Word8 -> IO r
f (ForeignPtr Word8 -> Int -> IO r
finalize ForeignPtr Word8
fp) Ptr Word8
buf
{-# INLINE createCPS #-}

{-
withBuffer
    :: Int
    -> (Ptr Word8 -> Int -> IO r)
    -> (r ->
withBuffer bufLen
{-# INLINE withBuffer #-}
-}

createAndTrimCPS
    :: Int
    -> (Ptr Word8 -> (Int -> IO ByteString) -> IO r)
    -> IO r
createAndTrimCPS :: forall r.
Int -> (Ptr Word8 -> (Int -> IO ByteString) -> IO r) -> IO r
createAndTrimCPS Int
maxLen Ptr Word8 -> (Int -> IO ByteString) -> IO r
action = Bool -> IO r -> IO r
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
maxLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO r -> IO r) -> IO r -> IO r
forall a b. (a -> b) -> a -> b
$ do
    ForeignPtr Word8
fp <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
B.mallocByteString Int
maxLen
    ForeignPtr Word8 -> (Ptr Word8 -> IO r) -> IO r
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
fp ((Ptr Word8 -> IO r) -> IO r) -> (Ptr Word8 -> IO r) -> IO r
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
p -> Ptr Word8 -> (Int -> IO ByteString) -> IO r
action Ptr Word8
p ((Int -> IO ByteString) -> IO r) -> (Int -> IO ByteString) -> IO r
forall a b. (a -> b) -> a -> b
$ \Int
len ->
        if   Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maxLen
        then Int -> (ForeignPtr Word8 -> IO ()) -> IO ByteString
createFp Int
len (\ForeignPtr Word8
fp' -> ForeignPtr Word8 -> ForeignPtr Word8 -> Int -> IO ()
memcpyFp ForeignPtr Word8
fp' ForeignPtr Word8
fp Int
len)
             -- ^ apparently @fp@ will get GCed automatically, up to GHC
        else ForeignPtr Word8 -> Int -> IO ByteString
B.mkDeferredByteString ForeignPtr Word8
fp Int
maxLen
{-# INLINE createAndTrimCPS #-}

unsafeCreateAndTrimCPS
    :: Int
    -> (Ptr Word8 -> (Int -> IO ByteString) -> IO r)
    -> r
unsafeCreateAndTrimCPS :: forall r. Int -> (Ptr Word8 -> (Int -> IO ByteString) -> IO r) -> r
unsafeCreateAndTrimCPS Int
l Ptr Word8 -> (Int -> IO ByteString) -> IO r
f =
    IO r -> r
forall a. IO a -> a
unsafeDupablePerformIOByteString (Int -> (Ptr Word8 -> (Int -> IO ByteString) -> IO r) -> IO r
forall r.
Int -> (Ptr Word8 -> (Int -> IO ByteString) -> IO r) -> IO r
createAndTrimCPS Int
l Ptr Word8 -> (Int -> IO ByteString) -> IO r
f)
{-# INLINE unsafeCreateAndTrimCPS #-}

createAndTrimFailable
    :: Int
    -> (Ptr Word8 -> IO (Either e Int))
    -> IO (Either e ByteString)
createAndTrimFailable :: forall e.
Int -> (Ptr Word8 -> IO (Either e Int)) -> IO (Either e ByteString)
createAndTrimFailable Int
l Ptr Word8 -> IO (Either e Int)
action = Int
-> (ForeignPtr Word8 -> IO (Either e Int))
-> IO (Either e ByteString)
forall e.
Int
-> (ForeignPtr Word8 -> IO (Either e Int))
-> IO (Either e ByteString)
createFpAndTrimFailable Int
l ((Ptr Word8 -> IO (Either e Int))
-> ForeignPtr Word8 -> IO (Either e Int)
forall res. (Ptr Word8 -> IO res) -> ForeignPtr Word8 -> IO res
wrapAction Ptr Word8 -> IO (Either e Int)
action)
{-# INLINE createAndTrimFailable #-}

-- TODO how do I omit the Either allocation?
createFpAndTrimFailable
    :: Int
    -> (ForeignPtr Word8 -> IO (Either e Int))
    -> IO (Either e ByteString)
createFpAndTrimFailable :: forall e.
Int
-> (ForeignPtr Word8 -> IO (Either e Int))
-> IO (Either e ByteString)
createFpAndTrimFailable Int
maxLen ForeignPtr Word8 -> IO (Either e Int)
action = Bool -> IO (Either e ByteString) -> IO (Either e ByteString)
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
maxLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO (Either e ByteString) -> IO (Either e ByteString))
-> IO (Either e ByteString) -> IO (Either e ByteString)
forall a b. (a -> b) -> a -> b
$ do
    ForeignPtr Word8
fp <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
B.mallocByteString Int
maxLen
    ForeignPtr Word8 -> IO (Either e Int)
action ForeignPtr Word8
fp IO (Either e Int)
-> (Either e Int -> IO (Either e ByteString))
-> IO (Either e ByteString)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Right Int
len ->
        if   Int
len Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
maxLen
        then ByteString -> Either e ByteString
forall a b. b -> Either a b
Right (ByteString -> Either e ByteString)
-> IO ByteString -> IO (Either e ByteString)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> (ForeignPtr Word8 -> IO ()) -> IO ByteString
createFp Int
len (\ForeignPtr Word8
fp' -> ForeignPtr Word8 -> ForeignPtr Word8 -> Int -> IO ()
memcpyFp ForeignPtr Word8
fp' ForeignPtr Word8
fp Int
len)
             -- ^ apparently @fp@ will get GCed automatically, up to GHC
        else ByteString -> Either e ByteString
forall a b. b -> Either a b
Right (ByteString -> Either e ByteString)
-> IO ByteString -> IO (Either e ByteString)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ForeignPtr Word8 -> Int -> IO ByteString
B.mkDeferredByteString ForeignPtr Word8
fp Int
maxLen
      Left  e
err -> Either e ByteString -> IO (Either e ByteString)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Either e ByteString -> IO (Either e ByteString))
-> Either e ByteString -> IO (Either e ByteString)
forall a b. (a -> b) -> a -> b
$ e -> Either e ByteString
forall a b. a -> Either a b
Left e
err
{-# INLINE createFpAndTrimFailable #-}

createUptoNFailable
    :: Int
    -> (Ptr Word8 -> IO (Either e Int))
    -> IO (Either e ByteString)
createUptoNFailable :: forall e.
Int -> (Ptr Word8 -> IO (Either e Int)) -> IO (Either e ByteString)
createUptoNFailable Int
l Ptr Word8 -> IO (Either e Int)
action = Int
-> (ForeignPtr Word8 -> IO (Either e Int))
-> IO (Either e ByteString)
forall e.
Int
-> (ForeignPtr Word8 -> IO (Either e Int))
-> IO (Either e ByteString)
createFpUptoNFailable Int
l ((Ptr Word8 -> IO (Either e Int))
-> ForeignPtr Word8 -> IO (Either e Int)
forall res. (Ptr Word8 -> IO res) -> ForeignPtr Word8 -> IO res
wrapAction Ptr Word8 -> IO (Either e Int)
action)
{-# INLINE createUptoNFailable #-}

createFpUptoNFailable
    :: Int
    -> (ForeignPtr Word8 -> IO (Either e Int))
    -> IO (Either e ByteString)
createFpUptoNFailable :: forall e.
Int
-> (ForeignPtr Word8 -> IO (Either e Int))
-> IO (Either e ByteString)
createFpUptoNFailable Int
maxLen ForeignPtr Word8 -> IO (Either e Int)
action = Bool -> IO (Either e ByteString) -> IO (Either e ByteString)
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
maxLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO (Either e ByteString) -> IO (Either e ByteString))
-> IO (Either e ByteString) -> IO (Either e ByteString)
forall a b. (a -> b) -> a -> b
$ do
    ForeignPtr Word8
fp <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
B.mallocByteString Int
maxLen
    ForeignPtr Word8 -> IO (Either e Int)
action ForeignPtr Word8
fp IO (Either e Int)
-> (Either e Int -> IO (Either e ByteString))
-> IO (Either e ByteString)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Right Int
len -> ByteString -> Either e ByteString
forall a b. b -> Either a b
Right (ByteString -> Either e ByteString)
-> IO ByteString -> IO (Either e ByteString)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ForeignPtr Word8 -> Int -> IO ByteString
B.mkDeferredByteString ForeignPtr Word8
fp Int
len
      Left  e
err -> Either e ByteString -> IO (Either e ByteString)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Either e ByteString -> IO (Either e ByteString))
-> Either e ByteString -> IO (Either e ByteString)
forall a b. (a -> b) -> a -> b
$ e -> Either e ByteString
forall a b. a -> Either a b
Left e
err
{-# INLINE createFpUptoNFailable #-}

createFailable
    :: Int
    -> (Ptr Word8 -> IO (Either e Int))
    -> IO (Either e ByteString)
createFailable :: forall e.
Int -> (Ptr Word8 -> IO (Either e Int)) -> IO (Either e ByteString)
createFailable Int
l Ptr Word8 -> IO (Either e Int)
action = Int
-> (ForeignPtr Word8 -> IO (Either e Int))
-> IO (Either e ByteString)
forall e.
Int
-> (ForeignPtr Word8 -> IO (Either e Int))
-> IO (Either e ByteString)
createFpFailable Int
l ((Ptr Word8 -> IO (Either e Int))
-> ForeignPtr Word8 -> IO (Either e Int)
forall res. (Ptr Word8 -> IO res) -> ForeignPtr Word8 -> IO res
wrapAction Ptr Word8 -> IO (Either e Int)
action)
{-# INLINE createFailable #-}

-- TODO how do I omit the Either allocation?
createFpFailable
    :: Int
    -> (ForeignPtr Word8 -> IO (Either e Int))
    -> IO (Either e ByteString)
createFpFailable :: forall e.
Int
-> (ForeignPtr Word8 -> IO (Either e Int))
-> IO (Either e ByteString)
createFpFailable Int
maxLen ForeignPtr Word8 -> IO (Either e Int)
action = Bool -> IO (Either e ByteString) -> IO (Either e ByteString)
forall a. (?callStack::CallStack) => Bool -> a -> a
assert (Int
maxLen Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
0) (IO (Either e ByteString) -> IO (Either e ByteString))
-> IO (Either e ByteString) -> IO (Either e ByteString)
forall a b. (a -> b) -> a -> b
$ do
    ForeignPtr Word8
fp <- Int -> IO (ForeignPtr Word8)
forall a. Int -> IO (ForeignPtr a)
B.mallocByteString Int
maxLen
    ForeignPtr Word8 -> IO (Either e Int)
action ForeignPtr Word8
fp IO (Either e Int)
-> (Either e Int -> IO (Either e ByteString))
-> IO (Either e ByteString)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: Type -> Type) a b. Monad m => m a -> (a -> m b) -> m b
>>= \case
      Right Int
len ->
        -- TODO does not check for correctness (len <= maxLen)!! don't lie!!!!
        ByteString -> Either e ByteString
forall a b. b -> Either a b
Right (ByteString -> Either e ByteString)
-> IO ByteString -> IO (Either e ByteString)
forall (f :: Type -> Type) a b. Functor f => (a -> b) -> f a -> f b
<$> ForeignPtr Word8 -> Int -> IO ByteString
B.mkDeferredByteString ForeignPtr Word8
fp Int
len
      Left  e
err -> Either e ByteString -> IO (Either e ByteString)
forall a. a -> IO a
forall (f :: Type -> Type) a. Applicative f => a -> f a
pure (Either e ByteString -> IO (Either e ByteString))
-> Either e ByteString -> IO (Either e ByteString)
forall a b. (a -> b) -> a -> b
$ e -> Either e ByteString
forall a b. a -> Either a b
Left e
err
{-# INLINE createFpFailable #-}

-- TODO probably don't export
wrapAction :: (Ptr Word8 -> IO res) -> ForeignPtr Word8 -> IO res
wrapAction :: forall res. (Ptr Word8 -> IO res) -> ForeignPtr Word8 -> IO res
wrapAction = (ForeignPtr Word8 -> (Ptr Word8 -> IO res) -> IO res)
-> (Ptr Word8 -> IO res) -> ForeignPtr Word8 -> IO res
forall a b c. (a -> b -> c) -> b -> a -> c
flip ForeignPtr Word8 -> (Ptr Word8 -> IO res) -> IO res
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr
  -- Cannot use unsafeWithForeignPtr, because action can diverge

unsafeDupablePerformIOByteString :: IO a -> a
-- Why does this exist? In base-4.15.1.0 until at least base-4.18.0.0,
-- the version of unsafeDupablePerformIO in base prevents unboxing of
-- its results with an opaque call to GHC.Exts.lazy, for reasons described
-- in Note [unsafePerformIO and strictness] in GHC.IO.Unsafe. (See
-- https://hackage.haskell.org/package/base-4.18.0.0/docs/src/GHC.IO.Unsafe.html#line-30 .)
-- Even if we accept the (very questionable) premise that the sort of
-- function described in that note should work, we expect no such
-- calls to be made in the context of bytestring.  (And we really want
-- unboxing!)
unsafeDupablePerformIOByteString :: forall a. IO a -> a
unsafeDupablePerformIOByteString (IO State# RealWorld -> (# State# RealWorld, a #)
act) =
    case (State# RealWorld -> (# State# RealWorld, a #))
-> (# State# RealWorld, a #)
forall o. (State# RealWorld -> o) -> o
runRW# State# RealWorld -> (# State# RealWorld, a #)
act of (# State# RealWorld
_, a
res #) -> a
res