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