{-# LANGUAGE AllowAmbiguousTypes #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE DataKinds #-}
{-# LANGUAGE DuplicateRecordFields #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedRecordDot #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE TypeOperators #-}
{-# LANGUAGE UndecidableInstances #-}
{-# LANGUAGE NoFieldSelectors #-}

module Data.Persist.Internal
  ( (:!:) (..)

    -- * The Get type
  , Get (..)
  , GetEnv (..)
  , GetException (..)
  , getOffset
  , failGet
  , runGet
  , runGetIO
  , unsafeGetPrefix

    -- * The Put type
  , Put (..)
  , PutEnv (..)
  , PutException (..)
  , Chunk (..)
  , evalPut
  , evalPutStrictIO
  , evalPutLazy
  , evalPutLazyIO
  , grow

    -- * Size reservations
  , PutSize (..)
  ) where

import Control.Exception
import Control.Monad
import qualified Control.Monad.Fail as Fail
import Data.ByteString (ByteString)
import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as B
import qualified Data.ByteString.Lazy.Internal as BL
#if MIN_VERSION_base(4,20,0)
import Data.Foldable (foldlM)
#else
import Data.Foldable (foldl', foldlM)
#endif
import Data.IORef
import Data.List.NonEmpty (NonEmpty (..))
import Data.Word
import Foreign
  ( ForeignPtr
  , Ptr
  , allocaBytes
  , finalizerFree
  , free
  , mallocBytes
  , minusPtr
  , newForeignPtr
  , plusPtr
  , reallocBytes
  , withForeignPtr
  )
import Foreign.Marshal.Utils (copyBytes)
import System.IO.Unsafe

#include "MachDeps.h"

data a :!: b = !a :!: !b
infixl 2 :!:

data GetEnv = GetEnv
  { GetEnv -> ForeignPtr Word8
buf :: !(ForeignPtr Word8)
  , GetEnv -> Ptr Word8
begin :: {-# UNPACK #-} !(Ptr Word8)
  , GetEnv -> Ptr Word8
end :: {-# UNPACK #-} !(Ptr Word8)
  , GetEnv -> Ptr Word8
tmp :: {-# UNPACK #-} !(Ptr Word8)
  }

newtype Get a = Get
  { forall a. Get a -> GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)
unGet :: GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)
  }

instance Functor Get where
  fmap :: forall a b. (a -> b) -> Get a -> Get b
fmap a -> b
f Get a
m = (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Get b
forall a. (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
Get ((GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Get b)
-> (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Get b
forall a b. (a -> b) -> a -> b
$ \GetEnv
e Ptr Word8
p -> do
    Ptr Word8
p' :!: a
x <- Get a
m.unGet GetEnv
e Ptr Word8
p
    (Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b))
-> (Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b)
forall a b. (a -> b) -> a -> b
$! Ptr Word8
p' Ptr Word8 -> b -> Ptr Word8 :!: b
forall a b. a -> b -> a :!: b
:!: a -> b
f a
x
  {-# INLINE fmap #-}

instance Applicative Get where
  pure :: forall a. a -> Get a
pure a
a = (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
forall a. (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
Get ((GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a)
-> (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
forall a b. (a -> b) -> a -> b
$ \GetEnv
_ Ptr Word8
p -> (Ptr Word8 :!: a) -> IO (Ptr Word8 :!: a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Ptr Word8 :!: a) -> IO (Ptr Word8 :!: a))
-> (Ptr Word8 :!: a) -> IO (Ptr Word8 :!: a)
forall a b. (a -> b) -> a -> b
$! Ptr Word8
p Ptr Word8 -> a -> Ptr Word8 :!: a
forall a b. a -> b -> a :!: b
:!: a
a
  {-# INLINE pure #-}

  Get (a -> b)
f <*> :: forall a b. Get (a -> b) -> Get a -> Get b
<*> Get a
a = (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Get b
forall a. (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
Get ((GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Get b)
-> (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Get b
forall a b. (a -> b) -> a -> b
$ \GetEnv
e Ptr Word8
p -> do
    Ptr Word8
p' :!: a -> b
f' <- Get (a -> b)
f.unGet GetEnv
e Ptr Word8
p
    Ptr Word8
p'' :!: a
a' <- Get a
a.unGet GetEnv
e Ptr Word8
p'
    (Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b))
-> (Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b)
forall a b. (a -> b) -> a -> b
$! Ptr Word8
p'' Ptr Word8 -> b -> Ptr Word8 :!: b
forall a b. a -> b -> a :!: b
:!: a -> b
f' a
a'
  {-# INLINE (<*>) #-}

  Get a
m1 *> :: forall a b. Get a -> Get b -> Get b
*> Get b
m2 = do
    Get a -> Get ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void Get a
m1
    Get b
m2
  {-# INLINE (*>) #-}

instance Monad Get where
  Get a
m >>= :: forall a b. Get a -> (a -> Get b) -> Get b
>>= a -> Get b
f = (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Get b
forall a. (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
Get ((GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Get b)
-> (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Get b
forall a b. (a -> b) -> a -> b
$ \GetEnv
e Ptr Word8
p -> do
    Ptr Word8
p' :!: a
x <- Get a
m.unGet GetEnv
e Ptr Word8
p
    (a -> Get b
f a
x).unGet GetEnv
e Ptr Word8
p'
  {-# INLINE (>>=) #-}

#if !MIN_VERSION_base(4,11,0)
  fail = Fail.fail
  {-# INLINE fail #-}
#endif

data GetException
  = LengthException Int String
  | CharException Int String
  | EOFException Int String
  | GenericGetException Int String
  deriving (GetException -> GetException -> Bool
(GetException -> GetException -> Bool)
-> (GetException -> GetException -> Bool) -> Eq GetException
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: GetException -> GetException -> Bool
== :: GetException -> GetException -> Bool
$c/= :: GetException -> GetException -> Bool
/= :: GetException -> GetException -> Bool
Eq, Int -> GetException -> ShowS
[GetException] -> ShowS
GetException -> String
(Int -> GetException -> ShowS)
-> (GetException -> String)
-> ([GetException] -> ShowS)
-> Show GetException
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> GetException -> ShowS
showsPrec :: Int -> GetException -> ShowS
$cshow :: GetException -> String
show :: GetException -> String
$cshowList :: [GetException] -> ShowS
showList :: [GetException] -> ShowS
Show)

instance Exception GetException

data PutException
  = PutSizeMissingStartChunk
  deriving (PutException -> PutException -> Bool
(PutException -> PutException -> Bool)
-> (PutException -> PutException -> Bool) -> Eq PutException
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: PutException -> PutException -> Bool
== :: PutException -> PutException -> Bool
$c/= :: PutException -> PutException -> Bool
/= :: PutException -> PutException -> Bool
Eq, Int -> PutException -> ShowS
[PutException] -> ShowS
PutException -> String
(Int -> PutException -> ShowS)
-> (PutException -> String)
-> ([PutException] -> ShowS)
-> Show PutException
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> PutException -> ShowS
showsPrec :: Int -> PutException -> ShowS
$cshow :: PutException -> String
show :: PutException -> String
$cshowList :: [PutException] -> ShowS
showList :: [PutException] -> ShowS
Show)

instance Exception PutException

instance Fail.MonadFail Get where
  fail :: forall a. String -> Get a
fail String
msg = (Int -> String -> GetException) -> String -> Get a
forall a. (Int -> String -> GetException) -> String -> Get a
failGet Int -> String -> GetException
GenericGetException (String
"Failed reading: " String -> ShowS
forall a. Semigroup a => a -> a -> a
<> String
msg)
  {-# INLINE fail #-}

getOffset :: Get Int
getOffset :: Get Int
getOffset = (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: Int)) -> Get Int
forall a. (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
Get ((GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: Int)) -> Get Int)
-> (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: Int)) -> Get Int
forall a b. (a -> b) -> a -> b
$ \GetEnv
e Ptr Word8
p -> (Ptr Word8 :!: Int) -> IO (Ptr Word8 :!: Int)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Ptr Word8 :!: Int) -> IO (Ptr Word8 :!: Int))
-> (Ptr Word8 :!: Int) -> IO (Ptr Word8 :!: Int)
forall a b. (a -> b) -> a -> b
$! Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8 :!: Int
forall a b. a -> b -> a :!: b
:!: (Ptr Word8
p Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` GetEnv
e.begin)
{-# INLINE getOffset #-}

failGet :: (Int -> String -> GetException) -> String -> Get a
failGet :: forall a. (Int -> String -> GetException) -> String -> Get a
failGet Int -> String -> GetException
ctor String
msg = do
  Int
offset <- Get Int
getOffset
  (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
forall a. (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
Get ((GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a)
-> (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
forall a b. (a -> b) -> a -> b
$ \GetEnv
_ Ptr Word8
_ -> GetException -> IO (Ptr Word8 :!: a)
forall e a. Exception e => e -> IO a
throwIO (Int -> String -> GetException
ctor Int
offset String
msg)

runGetIO :: Get a -> ByteString -> IO a
runGetIO :: forall a. Get a -> ByteString -> IO a
runGetIO Get a
m ByteString
s = IO a
run
 where
  run :: IO a
run = ForeignPtr Word8 -> (Ptr Word8 -> IO a) -> IO a
forall a b. ForeignPtr a -> (Ptr a -> IO b) -> IO b
withForeignPtr ForeignPtr Word8
buf ((Ptr Word8 -> IO a) -> IO a) -> (Ptr Word8 -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
p -> Int -> (Ptr Word8 -> IO a) -> IO a
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
8 ((Ptr Word8 -> IO a) -> IO a) -> (Ptr Word8 -> IO a) -> IO a
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
t -> do
    let env :: GetEnv
env = GetEnv {ForeignPtr Word8
buf :: ForeignPtr Word8
buf :: ForeignPtr Word8
buf, begin :: Ptr Word8
begin = Ptr Word8
p, end :: Ptr Word8
end = Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` (Int
pos Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Int
len), tmp :: Ptr Word8
tmp = Ptr Word8
t}
    Ptr Word8
_ :!: a
r <- Get a
m.unGet GetEnv
env (Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
pos)
    a -> IO a
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure a
r
  (B.PS ForeignPtr Word8
buf Int
pos Int
len) = ByteString
s

-- | Run the Get monad applies a 'get'-based parser on the input ByteString
runGet :: Get a -> ByteString -> Either String a
runGet :: forall a. Get a -> ByteString -> Either String a
runGet Get a
m ByteString
s = IO (Either String a) -> Either String a
forall a. IO a -> a
unsafePerformIO (IO (Either String a) -> Either String a)
-> IO (Either String a) -> Either String a
forall a b. (a -> b) -> a -> b
$ IO (Either String a)
-> (GetException -> IO (Either String a)) -> IO (Either String a)
forall e a. Exception e => IO a -> (e -> IO a) -> IO a
catch (a -> Either String a
forall a b. b -> Either a b
Right (a -> Either String a) -> IO a -> IO (Either String a)
forall (m :: * -> *) a b. Monad m => (a -> b) -> m a -> m b
<$!> Get a -> ByteString -> IO a
forall a. Get a -> ByteString -> IO a
runGetIO Get a
m ByteString
s) GetException -> IO (Either String a)
forall {f :: * -> *} {b}.
Applicative f =>
GetException -> f (Either String b)
handler
 where
  handler :: GetException -> f (Either String b)
handler (GetException
e :: GetException) = Either String b -> f (Either String b)
forall a. a -> f a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Either String b -> f (Either String b))
-> Either String b -> f (Either String b)
forall a b. (a -> b) -> a -> b
$ String -> Either String b
forall a b. a -> Either a b
Left (String -> Either String b) -> String -> Either String b
forall a b. (a -> b) -> a -> b
$ GetException -> String
forall e. Exception e => e -> String
displayException GetException
e
{-# NOINLINE runGet #-}

unsafeGetPrefix :: Int -> Get a -> Get a
unsafeGetPrefix :: forall a. Int -> Get a -> Get a
unsafeGetPrefix Int
prefixLength Get a
baseGet = (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
forall a. (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
Get ((GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a)
-> (GetEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Get a
forall a b. (a -> b) -> a -> b
$ \GetEnv
env Ptr Word8
p -> do
  let p' :: Ptr Word8
p' = Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
prefixLength
      env' :: GetEnv
env' = (\GetEnv {Ptr Word8
ForeignPtr Word8
buf :: GetEnv -> ForeignPtr Word8
begin :: GetEnv -> Ptr Word8
end :: GetEnv -> Ptr Word8
tmp :: GetEnv -> Ptr Word8
buf :: ForeignPtr Word8
begin :: Ptr Word8
end :: Ptr Word8
tmp :: Ptr Word8
..} -> GetEnv {end :: Ptr Word8
end = Ptr Word8
p', Ptr Word8
ForeignPtr Word8
buf :: ForeignPtr Word8
begin :: Ptr Word8
tmp :: Ptr Word8
buf :: ForeignPtr Word8
begin :: Ptr Word8
tmp :: Ptr Word8
..}) GetEnv
env
  Ptr Word8
_ :!: a
r <- Get a
baseGet.unGet GetEnv
env' Ptr Word8
p
  (Ptr Word8 :!: a) -> IO (Ptr Word8 :!: a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Ptr Word8 :!: a) -> IO (Ptr Word8 :!: a))
-> (Ptr Word8 :!: a) -> IO (Ptr Word8 :!: a)
forall a b. (a -> b) -> a -> b
$ Ptr Word8
p' Ptr Word8 -> a -> Ptr Word8 :!: a
forall a b. a -> b -> a :!: b
:!: a
r
{-# INLINE unsafeGetPrefix #-}

data Chunk = Chunk
  { Chunk -> Ptr Word8
begin :: {-# UNPACK #-} !(Ptr Word8)
  , Chunk -> Ptr Word8
end :: {-# UNPACK #-} !(Ptr Word8)
  }

data PutEnv = PutEnv
  { PutEnv -> IORef (NonEmpty Chunk)
chunks :: !(IORef (NonEmpty Chunk))
  , PutEnv -> IORef (Ptr Word8)
end :: !(IORef (Ptr Word8))
  , PutEnv -> Ptr Word8
tmp :: {-# UNPACK #-} !(Ptr Word8)
  }

newtype Put a = Put
  {forall a. Put a -> PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)
unPut :: PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)}

instance Functor Put where
  fmap :: forall a b. (a -> b) -> Put a -> Put b
fmap a -> b
f Put a
m = (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Put b
forall a. (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Put a
Put ((PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Put b)
-> (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Put b
forall a b. (a -> b) -> a -> b
$ \PutEnv
e Ptr Word8
p -> do
    Ptr Word8
p' :!: a
x <- Put a
m.unPut PutEnv
e Ptr Word8
p
    (Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b))
-> (Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b)
forall a b. (a -> b) -> a -> b
$! Ptr Word8
p' Ptr Word8 -> b -> Ptr Word8 :!: b
forall a b. a -> b -> a :!: b
:!: a -> b
f a
x
  {-# INLINE fmap #-}

instance Applicative Put where
  pure :: forall a. a -> Put a
pure a
a = (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Put a
forall a. (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Put a
Put ((PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Put a)
-> (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Put a
forall a b. (a -> b) -> a -> b
$ \PutEnv
_ Ptr Word8
p -> (Ptr Word8 :!: a) -> IO (Ptr Word8 :!: a)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Ptr Word8 :!: a) -> IO (Ptr Word8 :!: a))
-> (Ptr Word8 :!: a) -> IO (Ptr Word8 :!: a)
forall a b. (a -> b) -> a -> b
$! Ptr Word8
p Ptr Word8 -> a -> Ptr Word8 :!: a
forall a b. a -> b -> a :!: b
:!: a
a
  {-# INLINE pure #-}

  Put (a -> b)
f <*> :: forall a b. Put (a -> b) -> Put a -> Put b
<*> Put a
a = (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Put b
forall a. (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Put a
Put ((PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Put b)
-> (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Put b
forall a b. (a -> b) -> a -> b
$ \PutEnv
e Ptr Word8
p -> do
    Ptr Word8
p' :!: a -> b
f' <- Put (a -> b)
f.unPut PutEnv
e Ptr Word8
p
    Ptr Word8
p'' :!: a
a' <- Put a
a.unPut PutEnv
e Ptr Word8
p'
    (Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b))
-> (Ptr Word8 :!: b) -> IO (Ptr Word8 :!: b)
forall a b. (a -> b) -> a -> b
$! Ptr Word8
p'' Ptr Word8 -> b -> Ptr Word8 :!: b
forall a b. a -> b -> a :!: b
:!: a -> b
f' a
a'
  {-# INLINE (<*>) #-}

  Put a
m1 *> :: forall a b. Put a -> Put b -> Put b
*> Put b
m2 = do
    Put a -> Put ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void Put a
m1
    Put b
m2
  {-# INLINE (*>) #-}

instance Monad Put where
  Put a
m >>= :: forall a b. Put a -> (a -> Put b) -> Put b
>>= a -> Put b
f = (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Put b
forall a. (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Put a
Put ((PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Put b)
-> (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: b)) -> Put b
forall a b. (a -> b) -> a -> b
$ \PutEnv
e Ptr Word8
p -> do
    Ptr Word8
p' :!: a
x <- Put a
m.unPut PutEnv
e Ptr Word8
p
    (a -> Put b
f a
x).unPut PutEnv
e Ptr Word8
p'
  {-# INLINE (>>=) #-}

data PutSize a = PutSize
  { forall a. PutSize a -> Ptr Word8
sizePtr :: !(Ptr Word8)
  , forall a. PutSize a -> Ptr Word8
sizeStart :: !(Ptr Word8)
  , forall a. PutSize a -> Ptr Word8
chunkStart :: !(Ptr Word8)
  }

minChunkSize :: Int
minChunkSize :: Int
minChunkSize = Int
0x10000
{-# INLINE minChunkSize #-}

newChunk :: Int -> IO Chunk
newChunk :: Int -> IO Chunk
newChunk Int
size = do
  let n :: Int
n = Int -> Int -> Int
forall a. Ord a => a -> a -> a
max Int
size Int
minChunkSize
  Ptr Word8
p <- Int -> IO (Ptr Word8)
forall a. Int -> IO (Ptr a)
mallocBytes Int
n
  Chunk -> IO Chunk
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Chunk -> IO Chunk) -> Chunk -> IO Chunk
forall a b. (a -> b) -> a -> b
$! Ptr Word8 -> Ptr Word8 -> Chunk
Chunk Ptr Word8
p (Ptr Word8 -> Chunk) -> Ptr Word8 -> Chunk
forall a b. (a -> b) -> a -> b
$ Ptr Word8
p Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
n
{-# INLINE newChunk #-}

-- | Ensure that @n@ bytes can be written.
grow :: Int -> Put ()
grow :: Int -> Put ()
grow Int
n
  | Int
n Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
< Int
0 = String -> Put ()
forall a. HasCallStack => String -> a
error String
"grow: negative length"
  | Bool
otherwise = (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: ())) -> Put ()
forall a. (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: a)) -> Put a
Put ((PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: ())) -> Put ())
-> (PutEnv -> Ptr Word8 -> IO (Ptr Word8 :!: ())) -> Put ()
forall a b. (a -> b) -> a -> b
$ \PutEnv
e Ptr Word8
p -> do
      Ptr Word8
end <- IORef (Ptr Word8) -> IO (Ptr Word8)
forall a. IORef a -> IO a
readIORef PutEnv
e.end
      if Ptr Word8
end Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
p Int -> Int -> Bool
forall a. Ord a => a -> a -> Bool
>= Int
n
        then
          (Ptr Word8 :!: ()) -> IO (Ptr Word8 :!: ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Ptr Word8 :!: ()) -> IO (Ptr Word8 :!: ()))
-> (Ptr Word8 :!: ()) -> IO (Ptr Word8 :!: ())
forall a b. (a -> b) -> a -> b
$! Ptr Word8
p Ptr Word8 -> () -> Ptr Word8 :!: ()
forall a b. a -> b -> a :!: b
:!: ()
        else
          PutEnv -> Ptr Word8 -> Int -> IO (Ptr Word8 :!: ())
doGrow PutEnv
e Ptr Word8
p Int
n
{-# INLINE grow #-}

doGrow :: PutEnv -> Ptr Word8 -> Int -> IO (Ptr Word8 :!: ())
doGrow :: PutEnv -> Ptr Word8 -> Int -> IO (Ptr Word8 :!: ())
doGrow PutEnv
e Ptr Word8
p Int
n = do
  Chunk
k <- Int -> IO Chunk
newChunk Int
n
  IORef (NonEmpty Chunk)
-> (NonEmpty Chunk -> NonEmpty Chunk) -> IO ()
forall a. IORef a -> (a -> a) -> IO ()
modifyIORef' PutEnv
e.chunks ((NonEmpty Chunk -> NonEmpty Chunk) -> IO ())
-> (NonEmpty Chunk -> NonEmpty Chunk) -> IO ()
forall a b. (a -> b) -> a -> b
$ \case
    (Chunk
c :| [Chunk]
cs) ->
      let !c' :: Chunk
c' = (\Chunk {Ptr Word8
begin :: Chunk -> Ptr Word8
end :: Chunk -> Ptr Word8
begin :: Ptr Word8
end :: Ptr Word8
..} -> Chunk {end :: Ptr Word8
end = Ptr Word8
p, Ptr Word8
begin :: Ptr Word8
begin :: Ptr Word8
..}) Chunk
c
       in Chunk
k Chunk -> [Chunk] -> NonEmpty Chunk
forall a. a -> [a] -> NonEmpty a
:| Chunk
c' Chunk -> [Chunk] -> [Chunk]
forall a. a -> [a] -> [a]
: [Chunk]
cs
  IORef (Ptr Word8) -> Ptr Word8 -> IO ()
forall a. IORef a -> a -> IO ()
writeIORef PutEnv
e.end (Ptr Word8 -> IO ()) -> Ptr Word8 -> IO ()
forall a b. (a -> b) -> a -> b
$! Chunk
k.end
  (Ptr Word8 :!: ()) -> IO (Ptr Word8 :!: ())
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ((Ptr Word8 :!: ()) -> IO (Ptr Word8 :!: ()))
-> (Ptr Word8 :!: ()) -> IO (Ptr Word8 :!: ())
forall a b. (a -> b) -> a -> b
$! Chunk
k.begin Ptr Word8 -> () -> Ptr Word8 :!: ()
forall a b. a -> b -> a :!: b
:!: ()
{-# NOINLINE doGrow #-}

chunksLength :: [Chunk] -> Int
chunksLength :: [Chunk] -> Int
chunksLength = (Int -> Chunk -> Int) -> Int -> [Chunk] -> Int
forall b a. (b -> a -> b) -> b -> [a] -> b
forall (t :: * -> *) b a.
Foldable t =>
(b -> a -> b) -> b -> t a -> b
foldl' (\Int
s Chunk
c -> Int
s Int -> Int -> Int
forall a. Num a => a -> a -> a
+ Chunk
c.end Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Chunk
c.begin) Int
0
{-# INLINE chunksLength #-}

catChunks :: [Chunk] -> IO ByteString
catChunks :: [Chunk] -> IO ByteString
catChunks [Chunk]
chks = Int -> (Ptr Word8 -> IO ()) -> IO ByteString
B.create ([Chunk] -> Int
chunksLength [Chunk]
chks) ((Ptr Word8 -> IO ()) -> IO ByteString)
-> (Ptr Word8 -> IO ()) -> IO ByteString
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
p ->
  IO (Ptr Word8) -> IO ()
forall (f :: * -> *) a. Functor f => f a -> f ()
void
    (IO (Ptr Word8) -> IO ()) -> IO (Ptr Word8) -> IO ()
forall a b. (a -> b) -> a -> b
$ (Ptr Word8 -> Chunk -> IO (Ptr Word8))
-> Ptr Word8 -> [Chunk] -> IO (Ptr Word8)
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM
      ( \Ptr Word8
q Chunk
c -> do
          let n :: Int
n = Chunk
c.end Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Chunk
c.begin
          Ptr Word8 -> Ptr Word8 -> Int -> IO ()
forall a. Ptr a -> Ptr a -> Int -> IO ()
copyBytes Ptr Word8
q Chunk
c.begin Int
n
          Ptr Word8 -> IO ()
forall a. Ptr a -> IO ()
free Chunk
c.begin
          Ptr Word8 -> IO (Ptr Word8)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Ptr Word8
q Ptr Word8 -> Int -> Ptr Word8
forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
n)
      )
      Ptr Word8
p
    ([Chunk] -> IO (Ptr Word8)) -> [Chunk] -> IO (Ptr Word8)
forall a b. (a -> b) -> a -> b
$ [Chunk] -> [Chunk]
forall a. [a] -> [a]
reverse [Chunk]
chks
{-# INLINE catChunks #-}

evalPutIO :: Put a -> (a -> NonEmpty Chunk -> IO (a, b)) -> IO (a, b)
evalPutIO :: forall a b.
Put a -> (a -> NonEmpty Chunk -> IO (a, b)) -> IO (a, b)
evalPutIO Put a
p a -> NonEmpty Chunk -> IO (a, b)
chunkConsumer = do
  Chunk
k <- Int -> IO Chunk
newChunk Int
0
  IORef (NonEmpty Chunk)
chunks <- NonEmpty Chunk -> IO (IORef (NonEmpty Chunk))
forall a. a -> IO (IORef a)
newIORef (Chunk
k Chunk -> [Chunk] -> NonEmpty Chunk
forall a. a -> [a] -> NonEmpty a
:| [])
  IORef (Ptr Word8)
curEnd <- Ptr Word8 -> IO (IORef (Ptr Word8))
forall a. a -> IO (IORef a)
newIORef Chunk
k.end
  Ptr Word8
p' :!: a
r <- Int -> (Ptr Word8 -> IO (Ptr Word8 :!: a)) -> IO (Ptr Word8 :!: a)
forall a b. Int -> (Ptr a -> IO b) -> IO b
allocaBytes Int
8 ((Ptr Word8 -> IO (Ptr Word8 :!: a)) -> IO (Ptr Word8 :!: a))
-> (Ptr Word8 -> IO (Ptr Word8 :!: a)) -> IO (Ptr Word8 :!: a)
forall a b. (a -> b) -> a -> b
$ \Ptr Word8
tmp ->
    Put a
p.unPut PutEnv {IORef (NonEmpty Chunk)
chunks :: IORef (NonEmpty Chunk)
chunks :: IORef (NonEmpty Chunk)
chunks, end :: IORef (Ptr Word8)
end = IORef (Ptr Word8)
curEnd, Ptr Word8
tmp :: Ptr Word8
tmp :: Ptr Word8
tmp} Chunk
k.begin
  NonEmpty Chunk
cs <- IORef (NonEmpty Chunk) -> IO (NonEmpty Chunk)
forall a. IORef a -> IO a
readIORef IORef (NonEmpty Chunk)
chunks
  case NonEmpty Chunk
cs of
    (Chunk
x :| [Chunk]
xs) -> do
      let !x' :: Chunk
x' = (\Chunk {Ptr Word8
begin :: Chunk -> Ptr Word8
end :: Chunk -> Ptr Word8
begin :: Ptr Word8
end :: Ptr Word8
..} -> Chunk {end :: Ptr Word8
end = Ptr Word8
p', Ptr Word8
begin :: Ptr Word8
begin :: Ptr Word8
..}) Chunk
x
      a -> NonEmpty Chunk -> IO (a, b)
chunkConsumer a
r (Chunk
x' Chunk -> [Chunk] -> NonEmpty Chunk
forall a. a -> [a] -> NonEmpty a
:| [Chunk]
xs)
{-# INLINE evalPutIO #-}

evalPutStrictIO :: Put a -> IO (a, ByteString)
evalPutStrictIO :: forall a. Put a -> IO (a, ByteString)
evalPutStrictIO Put a
p = Put a
-> (a -> NonEmpty Chunk -> IO (a, ByteString))
-> IO (a, ByteString)
forall a b.
Put a -> (a -> NonEmpty Chunk -> IO (a, b)) -> IO (a, b)
evalPutIO Put a
p a -> NonEmpty Chunk -> IO (a, ByteString)
forall {a}. a -> NonEmpty Chunk -> IO (a, ByteString)
chunkHandler
 where
  chunkHandler :: a -> NonEmpty Chunk -> IO (a, ByteString)
chunkHandler a
r NonEmpty Chunk
cs = do
    ByteString
s <- case NonEmpty Chunk
cs of
      (Chunk
x :| []) -> Chunk -> IO ByteString
singleChunk Chunk
x
      (Chunk
x :| [Chunk]
xs) -> [Chunk] -> IO ByteString
catChunks (Chunk
x Chunk -> [Chunk] -> [Chunk]
forall a. a -> [a] -> [a]
: [Chunk]
xs)
    (a, ByteString) -> IO (a, ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
r, ByteString
s)
  singleChunk :: Chunk -> IO ByteString
singleChunk Chunk {Ptr Word8
begin :: Chunk -> Ptr Word8
end :: Chunk -> Ptr Word8
begin :: Ptr Word8
end :: Ptr Word8
..} = do
    case Ptr Word8
end Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
begin of
      Int
0 -> do
        Ptr Word8 -> IO ()
forall a. Ptr a -> IO ()
free Ptr Word8
begin
        ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
B.empty
      Int
newSize -> do
        Ptr Word8
newPtr <- Ptr Word8 -> Int -> IO (Ptr Word8)
forall a. Ptr a -> Int -> IO (Ptr a)
reallocBytes Ptr Word8
begin Int
newSize
        ForeignPtr Word8
foreignNewPtr <- FinalizerPtr Word8 -> Ptr Word8 -> IO (ForeignPtr Word8)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr Word8
forall a. FinalizerPtr a
finalizerFree Ptr Word8
newPtr
        ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ForeignPtr Word8 -> Int -> ByteString
B.BS ForeignPtr Word8
foreignNewPtr Int
newSize
{-# INLINE evalPutStrictIO #-}

evalPutLazyIO :: Put a -> IO (a, BL.ByteString)
evalPutLazyIO :: forall a. Put a -> IO (a, ByteString)
evalPutLazyIO Put a
p = Put a
-> (a -> NonEmpty Chunk -> IO (a, ByteString))
-> IO (a, ByteString)
forall a b.
Put a -> (a -> NonEmpty Chunk -> IO (a, b)) -> IO (a, b)
evalPutIO Put a
p a -> NonEmpty Chunk -> IO (a, ByteString)
forall {a}. a -> NonEmpty Chunk -> IO (a, ByteString)
chunkHandler
 where
  chunkHandler :: a -> NonEmpty Chunk -> IO (a, ByteString)
chunkHandler a
r NonEmpty Chunk
cs = do
    ByteString
s <- case NonEmpty Chunk
cs of
      (Chunk
x :| [Chunk]
xs) -> (ByteString -> Chunk -> IO ByteString)
-> ByteString -> [Chunk] -> IO ByteString
forall (t :: * -> *) (m :: * -> *) b a.
(Foldable t, Monad m) =>
(b -> a -> m b) -> b -> t a -> m b
foldlM ByteString -> Chunk -> IO ByteString
makeLBSChunk ByteString
BL.Empty (Chunk
x Chunk -> [Chunk] -> [Chunk]
forall a. a -> [a] -> [a]
: [Chunk]
xs)
    (a, ByteString) -> IO (a, ByteString)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (a
r, ByteString
s)
  makeLBSChunk :: BL.ByteString -> Chunk -> IO BL.ByteString
  makeLBSChunk :: ByteString -> Chunk -> IO ByteString
makeLBSChunk ByteString
lbsTail Chunk {Ptr Word8
begin :: Chunk -> Ptr Word8
end :: Chunk -> Ptr Word8
begin :: Ptr Word8
end :: Ptr Word8
..} = do
    case Ptr Word8
end Ptr Word8 -> Ptr Word8 -> Int
forall a b. Ptr a -> Ptr b -> Int
`minusPtr` Ptr Word8
begin of
      Int
0 -> do
        Ptr Word8 -> IO ()
forall a. Ptr a -> IO ()
free Ptr Word8
begin
        ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ByteString
lbsTail
      Int
newSize -> do
        Ptr Word8
newPtr <- Ptr Word8 -> Int -> IO (Ptr Word8)
forall a. Ptr a -> Int -> IO (Ptr a)
reallocBytes Ptr Word8
begin Int
newSize
        ForeignPtr Word8
foreignNewPtr <- FinalizerPtr Word8 -> Ptr Word8 -> IO (ForeignPtr Word8)
forall a. FinalizerPtr a -> Ptr a -> IO (ForeignPtr a)
newForeignPtr FinalizerPtr Word8
forall a. FinalizerPtr a
finalizerFree Ptr Word8
newPtr
        let strictChunk :: ByteString
strictChunk = ForeignPtr Word8 -> Int -> ByteString
B.BS ForeignPtr Word8
foreignNewPtr Int
newSize
        ByteString -> IO ByteString
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ByteString -> IO ByteString) -> ByteString -> IO ByteString
forall a b. (a -> b) -> a -> b
$ ByteString -> ByteString -> ByteString
BL.Chunk ByteString
strictChunk ByteString
lbsTail
{-# INLINE evalPutLazyIO #-}

evalPut :: Put a -> (a, ByteString)
evalPut :: forall a. Put a -> (a, ByteString)
evalPut Put a
p = IO (a, ByteString) -> (a, ByteString)
forall a. IO a -> a
unsafePerformIO (IO (a, ByteString) -> (a, ByteString))
-> IO (a, ByteString) -> (a, ByteString)
forall a b. (a -> b) -> a -> b
$ Put a -> IO (a, ByteString)
forall a. Put a -> IO (a, ByteString)
evalPutStrictIO Put a
p
{-# NOINLINE evalPut #-}

evalPutLazy :: Put a -> (a, BL.ByteString)
evalPutLazy :: forall a. Put a -> (a, ByteString)
evalPutLazy Put a
p = IO (a, ByteString) -> (a, ByteString)
forall a. IO a -> a
unsafePerformIO (IO (a, ByteString) -> (a, ByteString))
-> IO (a, ByteString) -> (a, ByteString)
forall a b. (a -> b) -> a -> b
$ Put a -> IO (a, ByteString)
forall a. Put a -> IO (a, ByteString)
evalPutLazyIO Put a
p
{-# NOINLINE evalPutLazy #-}