{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE MagicHash #-}
{-# LANGUAGE UnboxedTuples #-}
module Proto3.Wire.Reverse.Internal
( BuildR(..)
, appendBuildR
, foldlRVector
, toBuildR
, fromBuildR
, etaBuildR
, runBuildR
, withUnused
, withTotal
, withLengthOf
, withLengthOf#
, reallocate
, prependChunk
, ReverseChunks(..)
, prependReverseChunks
, ensure
, ensure#
, unsafeConsume
, floatToWord32
, doubleToWord64
) where
import Control.Exception ( bracket )
import Control.Monad.Trans.State.Strict ( State, runState, state )
import qualified Data.ByteString as B
import qualified Data.ByteString.Internal as BI
import qualified Data.ByteString.Builder.Extra as BB
import qualified Data.ByteString.Lazy as BL
import qualified Data.ByteString.Lazy.Internal as BLI
import Data.IORef ( IORef, newIORef,
readIORef, writeIORef )
import qualified Data.Primitive as P
import qualified Data.Vector.Generic as VG
import Data.Vector.Generic ( Vector )
import Data.Word ( Word8, Word32, Word64 )
import Foreign ( Storable(..),
castPtrToStablePtr,
castStablePtrToPtr,
freeStablePtr, newStablePtr,
deRefStablePtr )
import GHC.Exts ( Addr#, Int#, MutVar#,
RealWorld, StablePtr#, State#,
addrToAny#, int2Addr#,
touch# )
import GHC.ForeignPtr ( ForeignPtr(..),
ForeignPtrContents(..) )
import GHC.IO ( IO(..) )
import GHC.IORef ( IORef(..) )
import GHC.Int ( Int(..) )
import GHC.Ptr ( Ptr(..), plusPtr )
import GHC.Stable ( StablePtr(..) )
import GHC.STRef ( STRef(..) )
import System.IO.Unsafe ( unsafePerformIO )
#if MIN_VERSION_primitive(0,7,0)
#define PTR P.Ptr
#else
#define PTR P.Addr
#endif
newtype BuildR = BuildR
(Addr# -> Int# -> State# RealWorld -> (# Addr#, Int#, State# RealWorld #))
instance Semigroup BuildR
where
(<>) = appendBuildR
{-# INLINE (<>) #-}
instance Monoid BuildR
where
mempty = BuildR (\v u s -> (# v, u, s #))
{-# INLINE mempty #-}
mappend = (<>)
{-# INLINE mappend #-}
instance Show BuildR
where
showsPrec prec builder =
showParen (prec > 10)
(showString "Proto3.Wire.Reverse.lazyByteString " . shows bytes)
where
bytes = snd (runBuildR builder)
appendBuildR :: BuildR -> BuildR -> BuildR
appendBuildR = \b c ->
let BuildR f = b
BuildR g = c
in
BuildR (\v0 u0 s0 -> case g v0 u0 s0 of (# v1, u1, s1 #) -> f v1 u1 s1)
{-# INLINE CONLIKE [1] appendBuildR #-}
foldlRVector :: Vector v a => (b -> a -> b) -> b -> v a -> b
foldlRVector f = \z v -> VG.foldr (flip f) z (VG.reverse v)
{-# INLINE foldlRVector #-}
toBuildR :: (Ptr Word8 -> Int -> IO (Ptr Word8, Int)) -> BuildR
toBuildR f =
BuildR $ \v0 u0 s0 ->
let IO g = f (Ptr v0) (I# u0) in
case g s0 of (# s1, (Ptr v1, I# u1) #) -> (# v1, u1, s1 #)
fromBuildR :: BuildR -> Ptr Word8 -> Int -> IO (Ptr Word8, Int)
fromBuildR (BuildR f) (Ptr v0) (I# u0) =
IO $ \s0 -> case f v0 u0 s0 of (# v1, u1, s1 #) -> (# s1, (Ptr v1, I# u1) #)
etaBuildR :: (a -> BuildR) -> a -> BuildR
etaBuildR f x = toBuildR $ \v u -> fromBuildR (f x) v u
data BuildRState = BuildRState
{ currentBuffer :: {-# UNPACK #-}!(P.MutableByteArray RealWorld)
, sealedBuffers :: BL.ByteString
}
allocateFields :: State (Int, Int) a -> (a, Int, Int)
allocateFields fields = (x, size, align)
where
(x, (off, align)) = runState fields (0, 1)
size = mod off align - off
allocatePrimitiveField :: Storable a => a -> State (Int, Int) Int
allocatePrimitiveField proxy = state $ \(prevOff, prevAlign) ->
let fieldWidth = sizeOf proxy
fieldAlign = alignment proxy
unaligned = prevOff - fieldWidth
nextOff = unaligned - mod unaligned fieldAlign
nextAlign = max prevAlign fieldAlign
in (nextOff, (nextOff, nextAlign))
scratchOffset, spaceOffset, stateOffset, metaDataSize, metaDataAlign :: Int
((scratchOffset, spaceOffset, stateOffset), metaDataSize, metaDataAlign) =
allocateFields $
(,,) <$> allocatePrimitiveField (undefined :: Double)
<*> allocatePrimitiveField (undefined :: Int)
<*> allocatePrimitiveField (undefined :: Ptr ())
smallChunkSize, defaultChunkSize :: Int
smallChunkSize = BB.smallChunkSize - metaDataSize
defaultChunkSize = BB.defaultChunkSize - metaDataSize
data MetaData
metaPtr :: Ptr Word8 -> Int -> Ptr MetaData
metaPtr v = plusPtr v . negate
readState :: Ptr MetaData -> IO (StablePtr (IORef BuildRState))
readState m = castPtrToStablePtr <$> peekByteOff m stateOffset
writeState :: Ptr MetaData -> StablePtr (IORef BuildRState) -> IO ()
writeState m = pokeByteOff m stateOffset . castStablePtrToPtr
readSpace :: Ptr MetaData -> IO Int
readSpace m = peekByteOff m spaceOffset
writeSpace :: Ptr MetaData -> Int -> IO ()
writeSpace m = pokeByteOff m spaceOffset
readTotal :: Ptr Word8 -> Int -> IO Int
readTotal v unused = do
space <- readSpace (metaPtr v unused)
let !total = space - unused
strictify total
strictify :: Int -> IO Int
strictify (I# x) = IO $ \s0 ->
case addrToAny# (int2Addr# x) of
(# y #) -> case touch# y s0 of
s1 -> (# s1, I# x #)
newBuffer ::
BL.ByteString ->
Int ->
IORef BuildRState ->
StablePtr (IORef BuildRState) ->
Int ->
IO (Ptr Word8)
newBuffer sealed (I# total) (IORef (STRef stateVar)) (StablePtr stateSP)
(I# unused) =
IO $ \s0 ->
case newBuffer# sealed total stateVar stateSP unused s0 of
(# s1, addr #) -> (# s1, Ptr addr #)
newBuffer# ::
BL.ByteString ->
Int# ->
MutVar# RealWorld BuildRState ->
StablePtr# (IORef BuildRState) ->
Int# ->
State# RealWorld ->
(# State# RealWorld, Addr# #)
newBuffer# sealed total stateVar stateSP unused s0 =
case go s0 of
(# s1, Ptr addr #) -> (# s1, addr #)
where
IO go = do
let allocation = metaDataSize + I# unused
buf <- P.newAlignedPinnedByteArray allocation metaDataAlign
let !(PTR base) = P.mutableByteArrayContents buf
!v = plusPtr (Ptr base) (metaDataSize + I# unused)
!m = plusPtr (Ptr base) metaDataSize
writeState m (StablePtr stateSP)
writeSpace m (I# unused + I# total)
let !nextState = BuildRState{currentBuffer = buf, sealedBuffers = sealed}
writeIORef (IORef (STRef stateVar)) nextState
pure v
data SealedState = SealedState
{ sealedSB :: BL.ByteString
, totalSB :: {-# UNPACK #-}!Int
, stateVarSB :: {-# UNPACK #-}!(IORef BuildRState)
, statePtrSB :: {-# UNPACK #-}!(StablePtr (IORef BuildRState))
, recycledSB :: Maybe (P.MutableByteArray RealWorld)
}
sealBuffer ::
Ptr Word8 ->
Int ->
IO SealedState
sealBuffer (Ptr addr) (I# u) = IO $ \s0 ->
case sealBuffer# addr u s0 of
(# s1, sealed, total, stateVar, statePtr, recycled #) ->
(# s1
, SealedState
{ sealedSB = sealed
, totalSB = I# total
, stateVarSB = IORef (STRef stateVar)
, statePtrSB = StablePtr statePtr
, recycledSB = recycled
}
#)
sealBuffer# ::
Addr# ->
Int# ->
State# RealWorld ->
(# State# RealWorld
, BL.ByteString
, Int#
, MutVar# RealWorld BuildRState
, StablePtr# (IORef BuildRState)
, Maybe (P.MutableByteArray RealWorld)
#)
sealBuffer# addr unused s0 =
case go s0 of
(# s1, (sealed, I# total, IORef (STRef sv), StablePtr sp, re) #) ->
(# s1, sealed, total, sv, sp, re #)
where
IO go = do
let v = Ptr addr
statePtr <- readState (metaPtr v (I# unused))
stateVar <- deRefStablePtr statePtr
BuildRState { currentBuffer = buffer, sealedBuffers = oldSealed } <-
readIORef stateVar
total <- readTotal v (I# unused)
let allocation = P.sizeofMutableByteArray buffer - metaDataSize
if allocation <= I# unused
then
pure (oldSealed, total, stateVar, statePtr, Just buffer)
else do
let !(PTR base) = P.mutableByteArrayContents buffer
!(P.MutableByteArray mba) = buffer
fp = ForeignPtr base (PlainPtr mba)
offset = metaDataSize + I# unused
finish trimmed recycled = do
let !newSealed = BLI.Chunk trimmed oldSealed
pure (newSealed, total, stateVar, statePtr, recycled)
untrimmed = BI.fromForeignPtr fp offset (allocation - I# unused)
if offset <= B.length untrimmed
then finish untrimmed Nothing
else finish (B.copy untrimmed) (Just buffer)
runBuildR :: BuildR -> (Int, BL.ByteString)
runBuildR f = unsafePerformIO $ do
stateVar <- newIORef undefined
bracket (newStablePtr stateVar) freeStablePtr $ \statePtr -> do
let u0 = smallChunkSize
v0 <- newBuffer BL.empty 0 stateVar statePtr u0
(v1, u1) <- fromBuildR f v0 u0
SealedState { sealedSB = bytes, totalSB = total } <- sealBuffer v1 u1
pure (total, bytes)
withUnused :: (Int -> BuildR) -> BuildR
withUnused f = toBuildR $ \v u -> fromBuildR (f u) v u
withTotal :: (Int -> BuildR) -> BuildR
withTotal f = withTotal# (\total -> f (I# total))
withTotal# :: (Int# -> BuildR) -> BuildR
withTotal# f = toBuildR $ \v u -> do
I# total <- readTotal v u
fromBuildR (f total) v u
withLengthOf :: (Int -> BuildR) -> BuildR -> BuildR
withLengthOf = \f g -> withLengthOf# (\len -> f (I# len)) g
{-# INLINE CONLIKE withLengthOf #-}
withLengthOf# :: (Int# -> BuildR) -> BuildR -> BuildR
withLengthOf# = \f g -> toBuildR $ \v0 u0 -> do
!before <- readTotal v0 u0
(v1, u1) <- fromBuildR g v0 u0
!after <- readTotal v1 u1
let !(I# len) = after - before
fromBuildR (f len) v1 u1
{-# INLINE CONLIKE [1] withLengthOf# #-}
reallocate :: Int -> BuildR
reallocate (I# required) = reallocate# required
reallocate# :: Int# -> BuildR
reallocate# required = toBuildR $ \v0 u0 -> do
SealedState
{ sealedSB = bytes
, totalSB = total
, stateVarSB = IORef (STRef stateVar)
, statePtrSB = StablePtr statePtr
} <- sealBuffer v0 u0
let !u1 = max (I# required) defaultChunkSize
v1 <- newBuffer bytes total (IORef (STRef stateVar)) (StablePtr statePtr) u1
pure (v1, u1)
{-# NOINLINE reallocate# #-}
afterPrependChunks :: SealedState -> IO (Ptr Word8, Int)
afterPrependChunks !st = IO $ \s0 ->
case afterPrependChunks# st s0 of
(# v1, u1, s1 #) -> (# s1, (Ptr v1, I# u1) #)
afterPrependChunks# ::
SealedState ->
State# RealWorld ->
(# Addr#, Int#, State# RealWorld #)
afterPrependChunks# SealedState
{ sealedSB = sealed
, totalSB = total
, stateVarSB = stateVar
, statePtrSB = statePtr
, recycledSB = recycled
} s0 =
case go s0 of (# s2, (Ptr v2, I# u2) #) -> (# v2, u2, s2 #)
where
IO go = case recycled of
Nothing -> do
let u1 = defaultChunkSize
v1 <- newBuffer sealed total stateVar statePtr u1
pure (v1, u1)
Just buf -> do
let u1 = P.sizeofMutableByteArray buf - metaDataSize
!(PTR base) = P.mutableByteArrayContents buf
!v1 = plusPtr (Ptr base) (metaDataSize + u1)
!m = plusPtr (Ptr base) metaDataSize
writeSpace m (u1 + total)
let !nextState = BuildRState
{ currentBuffer = buf, sealedBuffers = sealed }
writeIORef stateVar nextState
pure (v1, u1)
prependChunk :: B.ByteString -> BuildR
prependChunk (BI.PS (ForeignPtr ad ct) (I# off) (I# len))
| I# len == 0 = mempty
| otherwise = BuildR (\v u s -> prependChunk# v u s ad ct off len)
prependChunk# ::
Addr# ->
Int# ->
State# RealWorld ->
Addr# ->
ForeignPtrContents ->
Int# ->
Int# ->
(# Addr#, Int#, State# RealWorld #)
prependChunk# v u s ad ct off len = go v u s
where
BuildR go = toBuildR $ \v1 u1 -> do
SealedState
{ sealedSB = oldSealed
, totalSB = oldTotal
, stateVarSB = stateVar
, statePtrSB = statePtr
, recycledSB = recycled
} <- sealBuffer v1 u1
let chunk = BI.PS (ForeignPtr ad ct) (I# off) (I# len)
afterPrependChunks SealedState
{ sealedSB = BLI.Chunk chunk oldSealed
, totalSB = I# len + oldTotal
, stateVarSB = stateVar
, statePtrSB = statePtr
, recycledSB = recycled
}
newtype ReverseChunks = ReverseChunks { getReverseChunks :: BL.ByteString }
prependReverseChunks :: ReverseChunks -> BuildR
prependReverseChunks (ReverseChunks BLI.Empty) = mempty
prependReverseChunks
(ReverseChunks (BLI.Chunk (BI.PS (ForeignPtr ad ct) (I# off) (I# len)) cs)) =
BuildR (\v u s -> prependReverseChunks# v u s ad ct off len cs)
prependReverseChunks# ::
Addr# ->
Int# ->
State# RealWorld ->
Addr# ->
ForeignPtrContents ->
Int# ->
Int# ->
BL.ByteString ->
(# Addr#, Int#, State# RealWorld #)
prependReverseChunks# v0 u0 s0 ad ct off len cs0 = go v0 u0 s0
where
BuildR go = toBuildR $ \v1 u1 -> do
SealedState
{ sealedSB = oldSealed
, totalSB = oldTotal
, stateVarSB = stateVar
, statePtrSB = statePtr
, recycledSB = recycled
} <- sealBuffer v1 u1
let appendChunks !total sealed (BLI.Chunk c cs) =
appendChunks (B.length c + total) (BLI.Chunk c sealed) cs
appendChunks newTotal newSealed BLI.Empty =
afterPrependChunks SealedState
{ sealedSB = newSealed
, totalSB = newTotal
, stateVarSB = stateVar
, statePtrSB = statePtr
, recycledSB = recycled
}
let rchunks = BLI.Chunk (BI.PS (ForeignPtr ad ct) (I# off) (I# len)) cs0
appendChunks oldTotal oldSealed rchunks
ensure :: Int -> BuildR -> BuildR
ensure (I# required) f = ensure# required f
ensure# :: Int# -> BuildR -> BuildR
ensure# required (BuildR f) = BuildR $ \v u s ->
if I# required <= I# u
then f v u s
else let BuildR g = BuildR f <> reallocate# required in g v u s
unsafeConsume :: Int -> (Ptr Word8 -> IO ()) -> BuildR
unsafeConsume = \width f ->
toBuildR $ \v0 u0 -> do
let !m = - width
!v1 = plusPtr v0 m
!u1 = u0 + m
f v1
pure (v1, u1)
{-# INLINE unsafeConsume #-}
floatToWord32 :: Ptr Word8 -> Int -> Float -> IO Word32
floatToWord32 v u x = do
let m = metaPtr v u
pokeByteOff m scratchOffset x
peekByteOff m scratchOffset
doubleToWord64 :: Ptr Word8 -> Int -> Double -> IO Word64
doubleToWord64 v u x = do
let m = metaPtr v u
pokeByteOff m scratchOffset x
peekByteOff m scratchOffset