Tune to beat NewBinary
[packages/binary.git] / src / Data / Binary / Put.hs
1 -----------------------------------------------------------------------------
2 -- |
3 -- Module : Data.Binary.Put
4 -- Copyright : Lennart Kolmodin
5 -- License : BSD3-style (see LICENSE)
6 --
7 -- Maintainer : Lennart Kolmodin <kolmodin@dtek.chalmers.se>
8 -- Stability : stable
9 -- Portability : Portable to Hugs and GHC. Requires MPTCs
10 --
11 -- The Put monad. A monad for efficiently constructing lazy bytestrings.
12 --
13 -----------------------------------------------------------------------------
14
15 module Data.Binary.Put (
16
17 -- * The Put type
18 Put
19 , runPut
20
21 -- * Flushing the implicit parse state
22 , flush
23
24 -- * Primitives
25 , putWord8
26 , putByteString
27 , putLazyByteString
28
29 -- * Big-endian primitives
30 , putWord16be
31 , putWord32be
32 , putWord64be
33
34 -- * Little-endian primitives
35 , putWord16le
36 , putWord32le
37 , putWord64le
38
39 ) where
40
41 import Control.Monad.Writer
42
43 import Foreign
44
45 import Data.Monoid
46 import Data.Word
47 import Data.ByteString.Base (inlinePerformIO)
48 import qualified Data.ByteString.Base as S
49 import qualified Data.ByteString.Lazy as L
50
51 type Put = Writer Builder ()
52
53 -- | Run the 'Put' monad with a serialiser
54 runPut :: Put -> L.ByteString
55 runPut = runBuilder . execWriter
56
57 -- | Pop the ByteString we have constructed so far, if any, yielding a
58 -- new chunk in the result ByteString.
59 flush :: Put
60 flush = tell flushB
61
62 putWord8 :: Word8 -> Put
63 putWord8 = tell . singleton
64
65 -- | An efficient primitive to write a strict ByteString into the output buffer.
66 -- It flushes the current buffer, and writes the argument into a new chunk.
67 putByteString :: S.ByteString -> Put
68 putByteString = tell . putByteStringB
69
70 -- | Write a lazy ByteString efficiently, simply appending the lazy
71 -- ByteString chunks to the output buffer
72 putLazyByteString :: L.ByteString -> Put
73 putLazyByteString = tell . putLazyByteStringB
74
75 -- | Write a Word16 in big endian format
76 putWord16be :: Word16 -> Put
77 putWord16be = tell . putWord16beB
78 {-# INLINE putWord16be #-}
79
80 -- | Write a Word16 in little endian format
81 putWord16le :: Word16 -> Put
82 putWord16le = tell . putWord16leB
83
84 -- | Write a Word32 in big endian format
85 putWord32be :: Word32 -> Put
86 putWord32be = tell . putWord32beB
87
88 -- | Write a Word32 in little endian format
89 putWord32le :: Word32 -> Put
90 putWord32le = tell . putWord32leB
91
92 -- | Write a Word64 in big endian format
93 putWord64be :: Word64 -> Put
94 putWord64be = tell . putWord64beB
95
96 -- | Write a Word64 in little endian format
97 putWord64le :: Word64 -> Put
98 putWord64le = tell . putWord64leB
99
100 -- ---------------------------------------------------------------------
101 --
102 -- | The Builder monoid for efficiently constructing lazy bytestrings.
103 --
104
105 data Buffer = Buffer {-# UNPACK #-} !(ForeignPtr Word8)
106 {-# UNPACK #-} !Int -- offset
107 {-# UNPACK #-} !Int -- used bytes
108 {-# UNPACK #-} !Int -- length left
109
110 -- | The 'Builder' monoid abstracts over the construction of a lazy
111 -- bytestring by filling byte arrays piece by piece. As each buffer is
112 -- filled, it is \'popped\' off, to become a new chunk of the resulting
113 -- lazy 'L.ByteString'. All this is hidden from the user of the
114 -- 'Builder'.
115 --
116 -- Properties:
117 --
118 -- * @'runBuilder' 'empty' = 'L.empty'@
119 --
120 -- * @'runBuilder' ('append' x y) = 'L.append' ('runBuilder' x) ('runBuilder' y)@
121 --
122 -- * @'runBuilder' ('singleton' b) = 'L.singleton' b@
123 --
124 -- * @'runBuilder' ('putByteStringB' bs) = 'L.fromChunks' [bs]@
125 --
126 -- * @'runBuilder' ('putLazyByteStringB' bs) = bs@
127 --
128 newtype Builder = Builder {
129 unBuilder :: (Buffer -> [S.ByteString]) -> Buffer -> [S.ByteString]
130 }
131
132 instance Monoid Builder where
133 mempty = empty
134 mappend = append
135
136 empty :: Builder
137 empty = Builder id
138 {-# INLINE [1] empty #-}
139
140 append :: Builder -> Builder -> Builder
141 append (Builder f) (Builder g) = Builder (f . g)
142 {-# INLINE [1] append #-}
143
144 --
145 -- copied from Data.ByteString.Lazy
146 --
147 defaultSize :: Int
148 defaultSize = 32 * k - overhead
149 where k = 1024
150 overhead = 2 * sizeOf (undefined :: Int)
151
152 --
153 -- Run the builder monoid
154 --
155 runBuilder :: Builder -> L.ByteString
156 runBuilder m = S.LPS $ inlinePerformIO $ do
157 buf <- newBuffer defaultSize
158 return (unBuilder (m `append` flushB) (const []) buf)
159
160 -- | Sequence an IO operation on the buffer
161 unsafeLiftIO :: (Buffer -> IO Buffer) -> Builder
162 unsafeLiftIO f = Builder $ \ k buf -> inlinePerformIO $ do
163 buf' <- f buf
164 return (k buf')
165
166 -- | Get the size of the buffer
167 withSize :: (Int -> Builder) -> Builder
168 withSize f = Builder $ \ k buf@(Buffer _ _ _ l) ->
169 unBuilder (f l) k buf
170
171 -- | Map the resulting list of bytestrings.
172 mapBuilder :: ([S.ByteString] -> [S.ByteString]) -> Builder
173 mapBuilder f = Builder (f .)
174
175 -- | Pop the ByteString we have constructed so far, if any, yielding a
176 -- new chunk in the result ByteString.
177 flushB :: Builder
178 flushB = Builder $ \ k buf@(Buffer p o u l) ->
179 if u == 0
180 then k buf
181 else S.PS p o u : k (Buffer p (o+u) 0 l)
182 {-# INLINE [1] flush #-}
183
184 -- | Ensure that there are at least @n@ many bytes available.
185 ensureFree :: Int -> Builder
186 ensureFree n = n `seq` withSize $ \ l ->
187 if n <= l then empty else
188 flushB `append` unsafeLiftIO (const (newBuffer (max n defaultSize)))
189 {-# INLINE [1] ensureFree #-}
190
191 -- | Ensure that @n@ many bytes are available, and then use @f@ to write some
192 -- bytes into the memory.
193 writeN :: Int -> (Ptr Word8 -> IO ()) -> Builder
194 writeN n f = ensureFree n `append` unsafeLiftIO (writeNBuffer n f)
195 {-# INLINE [1] writeN #-}
196
197 writeNBuffer :: Int -> (Ptr Word8 -> IO ()) -> Buffer -> IO Buffer
198 writeNBuffer n f (Buffer fp o u l) = do
199 withForeignPtr fp (\p -> f (p `plusPtr` (o+u)))
200 return (Buffer fp o (u+n) (l-n))
201
202 newBuffer :: Int -> IO Buffer
203 newBuffer size = do
204 fp <- S.mallocByteString size
205 return $! Buffer fp 0 0 size
206
207 ------------------------------------------------------------------------
208
209 -- | Write a byte into the Builder's output buffer
210 singleton :: Word8 -> Builder
211 singleton = writeN 1 . flip poke
212 {-# INLINE putWord8 #-}
213
214 -- | Write a strict ByteString efficiently
215 putByteStringB :: S.ByteString -> Builder
216 putByteStringB bs = flushB `append` mapBuilder (bs :)
217
218 -- | Write a lazy ByteString efficiently
219 putLazyByteStringB :: L.ByteString -> Builder
220 putLazyByteStringB bs = flushB `append` mapBuilder (L.toChunks bs ++)
221
222 ------------------------------------------------------------------------
223
224 -- | Write a Word16 in big endian format
225 putWord16beB :: Word16 -> Builder
226 putWord16beB w16 =
227 let w1 = shiftR w16 8
228 w2 = w16 .&. 0xff
229 in
230 putWord8B (fromIntegral w1) `append`
231 putWord8B (fromIntegral w2)
232 {-# INLINE putWord16be #-}
233
234 -- | Write a Word16 in little endian format
235 putWord16leB :: Word16 -> Builder
236 -- putWord16leB w16 = writeN 2 (\p -> poke (castPtr p) w16)
237
238 putWord16leB w16 =
239 let w2 = shiftR w16 8
240 w1 = w16 .&. 0xff
241 in
242 singleton (fromIntegral w1) `append`
243 singleton (fromIntegral w2)
244 {-# INLINE putWord16le #-}
245
246 -- | Write a Word32 in big endian format
247 putWord32beB :: Word32 -> Builder
248 putWord32beB w32 =
249 let w1 = (w32 `shiftR` 24)
250 w2 = (w32 `shiftR` 16) .&. 0xff
251 w3 = (w32 `shiftR` 8) .&. 0xff
252 w4 = w32 .&. 0xff
253 in
254 singleton (fromIntegral w1) `append`
255 singleton (fromIntegral w2) `append`
256 singleton (fromIntegral w3) `append`
257 singleton (fromIntegral w4)
258 {-# INLINE putWord32be #-}
259
260 -- | Write a Word32 in little endian format
261 putWord32leB :: Word32 -> Builder
262 putWord32leB w32 =
263
264 -- on a little endian machine:
265 -- putWord32leB w32 = writeN 4 (\p -> poke (castPtr p) w32)
266
267 let w4 = (w32 `shiftR` 24)
268 w3 = (w32 `shiftR` 16) .&. 0xff
269 w2 = (w32 `shiftR` 8) .&. 0xff
270 w1 = w32 .&. 0xff
271 in
272 singleton (fromIntegral w1) `append`
273 singleton (fromIntegral w2) `append`
274 singleton (fromIntegral w3) `append`
275 singleton (fromIntegral w4)
276 {-# INLINE putWord32le #-}
277
278 -- | Write a Word64 in big endian format
279 putWord64beB :: Word64 -> Builder
280 putWord64beB w64 =
281 let w1 = shiftR w64 32
282 w2 = w64 .&. 0xffffffff
283 in
284 putWord32beB (fromIntegral w1) `append`
285 putWord32beB (fromIntegral w2)
286 {-# INLINE putWord64be #-}
287
288 -- | Write a Word64 in little endian format
289 putWord64leB :: Word64 -> Builder
290
291 -- on a little endian machine:
292 -- putWord64leB w64 = writeN 8 (\p -> poke (castPtr p) w64)
293
294 putWord64leB w64 =
295 let w2 = shiftR w64 32
296 w1 = w64 .&. 0xffffffff
297 in
298 putWord32leB (fromIntegral w1) `append`
299 putWord32leB (fromIntegral w2)
300 {-# INLINE putWord64le #-}
301
302 ------------------------------------------------------------------------
303 -- Some nice rules for Builder
304
305 {-# TRICKY RULES
306
307 "writeN/combine" forall s1 s2 f1 f2 .
308 bindP (writeN s1 f1) (writeN s2 f2) =
309 writeN (s1+s2) (\p -> f1 p >> f2 (p `plusPtr` s1))
310
311 "ensureFree/combine" forall a b .
312 bindP (ensureFree a) (ensureFree b) =
313 ensureFree (max a b)
314
315 "flush/combine"
316 bindP flush flush = flush
317
318 #-}