1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
|
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes #-}
-- | Convert a stream of bytestring @Builder@s into a stream of @ByteString@s.
--
-- Adapted from blaze-builder-enumerator, written by Michael Snoyman and Simon Meier.
--
-- Note that the functions here can work in any monad built on top of @IO@ or
-- @ST@.
--
-- Also provides @toByteStringIO*@ like "Blaze.ByteString.Builder"s, for
-- "Data.ByteString.Builder".
--
-- Since 0.1.9
--
module Data.Streaming.ByteString.Builder
( BuilderRecv
, BuilderPopper
, BuilderFinish
, newBuilderRecv
, newByteStringBuilderRecv
-- * toByteStringIO
, toByteStringIO
, toByteStringIOWith
, toByteStringIOWithBuffer
-- * Buffers
, Buffer
-- ** Status information
, freeSize
, sliceSize
, bufferSize
-- ** Creation and modification
, allocBuffer
, reuseBuffer
, nextSlice
-- ** Conversion to bytestings
, unsafeFreezeBuffer
, unsafeFreezeNonEmptyBuffer
-- * Buffer allocation strategies
, BufferAllocStrategy
, allNewBuffersStrategy
, reuseBufferStrategy
, defaultStrategy
)
where
import Control.Monad (when,unless)
import qualified Data.ByteString as S
import Data.ByteString.Builder (Builder)
import Data.ByteString.Builder.Extra (runBuilder, BufferWriter, Next(Done, More, Chunk))
import Data.ByteString.Internal (mallocByteString, ByteString(PS))
import Data.ByteString.Lazy.Internal (defaultChunkSize)
import Data.IORef (newIORef, writeIORef, readIORef)
import Data.Word (Word8)
import Foreign.ForeignPtr (ForeignPtr, withForeignPtr)
import Foreign.Ptr (plusPtr, minusPtr)
import Data.Streaming.ByteString.Builder.Buffer
-- | Provides a series of @ByteString@s until empty, at which point it provides
-- an empty @ByteString@.
--
-- Since 0.1.10.0
--
type BuilderPopper = IO S.ByteString
type BuilderRecv = Builder -> IO BuilderPopper
type BuilderFinish = IO (Maybe S.ByteString)
newBuilderRecv :: BufferAllocStrategy -> IO (BuilderRecv, BuilderFinish)
newBuilderRecv = newByteStringBuilderRecv
{-# INLINE newBuilderRecv #-}
newByteStringBuilderRecv :: BufferAllocStrategy -> IO (BuilderRecv, BuilderFinish)
newByteStringBuilderRecv (ioBufInit, nextBuf) = do
refBuf <- newIORef ioBufInit
return (push refBuf, finish refBuf)
where
finish refBuf = do
ioBuf <- readIORef refBuf
buf <- ioBuf
return $ unsafeFreezeNonEmptyBuffer buf
push refBuf builder = do
refWri <- newIORef $ Left $ runBuilder builder
return $ popper refBuf refWri
popper refBuf refWri = do
ioBuf <- readIORef refBuf
ebWri <- readIORef refWri
case ebWri of
Left bWri -> do
!buf@(Buffer _ _ op ope) <- ioBuf
(bytes, next) <- bWri op (ope `minusPtr` op)
let op' = op `plusPtr` bytes
case next of
Done -> do
writeIORef refBuf $ return $ updateEndOfSlice buf op'
return S.empty
More minSize bWri' -> do
let buf' = updateEndOfSlice buf op'
{-# INLINE cont #-}
cont mbs = do
-- sequencing the computation of the next buffer
-- construction here ensures that the reference to the
-- foreign pointer `fp` is lost as soon as possible.
ioBuf' <- nextBuf minSize buf'
writeIORef refBuf ioBuf'
writeIORef refWri $ Left bWri'
case mbs of
Just bs | not $ S.null bs -> return bs
_ -> popper refBuf refWri
cont $ unsafeFreezeNonEmptyBuffer buf'
Chunk bs bWri' -> do
let buf' = updateEndOfSlice buf op'
let yieldBS = do
nextBuf 1 buf' >>= writeIORef refBuf
writeIORef refWri $ Left bWri'
if S.null bs
then popper refBuf refWri
else return bs
case unsafeFreezeNonEmptyBuffer buf' of
Nothing -> yieldBS
Just bs' -> do
writeIORef refWri $ Right yieldBS
return bs'
Right action -> action
-- | Use a pre-existing buffer to 'toByteStringIOWith'.
--
-- Since 0.1.9
--
toByteStringIOWithBuffer :: Int
-> (ByteString -> IO ())
-> Builder
-> ForeignPtr Word8
-> IO ()
toByteStringIOWithBuffer initBufSize io b initBuf = do
go initBufSize initBuf (runBuilder b)
where
go bufSize buf = loop
where
loop :: BufferWriter -> IO ()
loop wr = do
(len, next) <- withForeignPtr buf (flip wr bufSize)
when (len > 0) (io $! PS buf 0 len)
case next of
Done -> return ()
More newBufSize nextWr
| newBufSize > bufSize -> do
newBuf <- mallocByteString newBufSize
go newBufSize newBuf nextWr
| otherwise -> loop nextWr
Chunk s nextWr -> do
unless (S.null s) (io s)
loop nextWr
-- | @toByteStringIOWith bufSize io b@ runs the builder @b@ with a buffer of
-- at least the size @bufSize@ and executes the 'IO' action @io@ whenever the
-- buffer is full.
--
-- Compared to 'toLazyByteStringWith' this function requires less allocation,
-- as the output buffer is only allocated once at the start of the
-- serialization and whenever something bigger than the current buffer size has
-- to be copied into the buffer, which should happen very seldomly for the
-- default buffer size of 32kb. Hence, the pressure on the garbage collector is
-- reduced, which can be an advantage when building long sequences of bytes.
--
-- Since 0.1.9
--
toByteStringIOWith :: Int -- ^ Buffer size (upper bounds
-- the number of bytes forced
-- per call to the 'IO' action).
-> (ByteString -> IO ()) -- ^ 'IO' action to execute per
-- full buffer, which is
-- referenced by a strict
-- 'S.ByteString'.
-> Builder -- ^ 'Builder' to run.
-> IO ()
toByteStringIOWith bufSize io b =
toByteStringIOWithBuffer bufSize io b =<< mallocByteString bufSize
{-# INLINE toByteStringIOWith #-}
-- | Run the builder with a 'defaultChunkSize'd buffer and execute the given
-- 'IO' action whenever the buffer is full or gets flushed.
--
-- @ 'toByteStringIO' = 'toByteStringIOWith' 'defaultChunkSize'@
--
-- Since 0.1.9
--
toByteStringIO :: (ByteString -> IO ())
-> Builder
-> IO ()
toByteStringIO = toByteStringIOWith defaultChunkSize
{-# INLINE toByteStringIO #-}
|
