1{-# LANGUAGE CPP, BangPatterns #-}
2
3-- |
4-- Module      : Blaze.ByteString.Builder.Internal.Poke
5-- Copyright   : (c) 2010 Simon Meier
6--               (c) 2010 Jasper van der Jeugt
7-- License     : BSD3-style (see LICENSE)
8--
9-- Maintainer  : Leon Smith <leon@melding-monads.com>
10-- Stability   : experimental
11-- Portability : tested on GHC only
12--
13-- A general and efficient write type that allows for the easy construction of
14-- builders for (smallish) bounded size writes to a buffer.
15--
16-- FIXME: Improve documentation.
17--
18module Blaze.ByteString.Builder.Internal.Write (
19  -- * Poking a buffer
20    Poke(..)
21  , pokeN
22
23  -- * Writing to abuffer
24  , Write(..)
25  , runWrite
26  , getBound
27  , getBound'
28  , getPoke
29
30  , exactWrite
31  , boundedWrite
32  , writeLiftIO
33  , writeIf
34  , writeEq
35  , writeOrdering
36  , writeOrd
37
38  -- * Constructing builders from writes
39  , fromWrite
40  , fromWriteSingleton
41  , fromWriteList
42
43  -- * Writing 'Storable's
44  , writeStorable
45  , fromStorable
46  , fromStorables
47
48  ) where
49
50import Foreign
51
52import qualified Data.Foldable as F
53import Control.Monad
54
55import Data.ByteString.Builder.Internal
56
57import Data.Monoid (Monoid(..))
58import Data.Semigroup (Semigroup(..))
59
60------------------------------------------------------------------------------
61-- Poking a buffer and writing to a buffer
62------------------------------------------------------------------------------
63
64-- Sadly GHC is not smart enough: code where we branch and each branch should
65-- execute a few IO actions and then return a value cannot be taught to GHC. At
66-- least not such that it returns the value of the branches unpacked.
67--
68-- Hmm.. at least he behaves much better for the Monoid instance of Write
69-- than the one for Poke. Serializing UTF-8 chars gets a slowdown of a
70-- factor 2 when 2 chars are composed. Perhaps I should try out the writeList
71-- instances also, as they may be more sensitive to to much work per Char.
72--
73
74-- | Changing a sequence of bytes starting from the given pointer. 'Poke's are
75-- the most primitive buffer manipulation. In most cases, you don't use the
76-- explicitely but as part of a 'Write', which also tells how many bytes will
77-- be changed at most.
78newtype Poke =
79    Poke { runPoke :: Ptr Word8 -> IO (Ptr Word8) }
80
81-- | A write of a bounded number of bytes.
82--
83-- When defining a function @write :: a -> Write@ for some @a@, then it is
84-- important to ensure that the bound on the number of bytes written is
85-- data-independent. Formally,
86--
87--  @ forall x y. getBound (write x) = getBound (write y) @
88--
89-- The idea is that this data-independent bound is specified such that the
90-- compiler can optimize the check, if there are enough free bytes in the buffer,
91-- to a single subtraction between the pointer to the next free byte and the
92-- pointer to the end of the buffer with this constant bound of the maximal
93-- number of bytes to be written.
94--
95data Write = Write {-# UNPACK #-} !Int Poke
96
97-- | Extract the 'Poke' action of a write.
98{-# INLINE getPoke #-}
99getPoke :: Write -> Poke
100getPoke (Write _ wio) = wio
101
102-- | Run the 'Poke' action of a write.
103{-# INLINE runWrite #-}
104runWrite :: Write -> Ptr Word8 -> IO (Ptr Word8)
105runWrite = runPoke . getPoke
106
107-- | Extract the maximal number of bytes that this write could write.
108{-# INLINE getBound #-}
109getBound :: Write -> Int
110getBound (Write bound _) = bound
111
112-- | Extract the maximal number of bytes that this write could write in any
113-- case. Assumes that the bound of the write is data-independent.
114{-# INLINE getBound' #-}
115getBound' :: String             -- ^ Name of caller: for debugging purposes.
116          -> (a -> Write)
117          -> Int
118getBound' msg write =
119    getBound $ write $ error $
120    "getBound' called from " ++ msg ++ ": write bound is not data-independent."
121
122instance Semigroup Poke where
123  {-# INLINE (<>) #-}
124  (Poke po1) <> (Poke po2) = Poke $ po1 >=> po2
125
126  {-# INLINE sconcat #-}
127  sconcat = F.foldr (<>) mempty
128
129instance Monoid Poke where
130  {-# INLINE mempty #-}
131  mempty = Poke $ return
132
133#if !(MIN_VERSION_base(4,11,0))
134  {-# INLINE mappend #-}
135  (Poke po1) `mappend` (Poke po2) = Poke $ po1 >=> po2
136
137  {-# INLINE mconcat #-}
138  mconcat = F.foldr mappend mempty
139#endif
140
141instance Semigroup Write where
142  {-# INLINE (<>) #-}
143  (Write bound1 w1) <> (Write bound2 w2) =
144    Write (bound1 + bound2) (w1 <> w2)
145
146  {-# INLINE sconcat #-}
147  sconcat = F.foldr (<>) mempty
148
149instance Monoid Write where
150  {-# INLINE mempty #-}
151  mempty = Write 0 mempty
152
153#if !(MIN_VERSION_base(4,11,0))
154  {-# INLINE mappend #-}
155  (Write bound1 w1) `mappend` (Write bound2 w2) =
156    Write (bound1 + bound2) (w1 `mappend` w2)
157
158  {-# INLINE mconcat #-}
159  mconcat = F.foldr mappend mempty
160#endif
161
162-- | @pokeN size io@ creates a write that denotes the writing of @size@ bytes
163-- to a buffer using the IO action @io@. Note that @io@ MUST write EXACTLY @size@
164-- bytes to the buffer!
165{-# INLINE pokeN #-}
166pokeN :: Int
167       -> (Ptr Word8 -> IO ()) -> Poke
168pokeN size io = Poke $ \op -> io op >> (return $! (op `plusPtr` size))
169
170
171-- | @exactWrite size io@ creates a bounded write that can later be converted to
172-- a builder that writes exactly @size@ bytes. Note that @io@ MUST write
173-- EXACTLY @size@ bytes to the buffer!
174{-# INLINE exactWrite #-}
175exactWrite :: Int
176           -> (Ptr Word8 -> IO ())
177           -> Write
178exactWrite size io = Write size (pokeN size io)
179
180-- | @boundedWrite size write@ creates a bounded write from a @write@ that does
181-- not write more than @size@ bytes.
182{-# INLINE boundedWrite #-}
183boundedWrite :: Int -> Poke -> Write
184boundedWrite = Write
185
186-- | @writeLiftIO io write@ creates a write executes the @io@ action to compute
187-- the value that is then written.
188{-# INLINE writeLiftIO #-}
189writeLiftIO :: (a -> Write) -> IO a -> Write
190writeLiftIO write io =
191    Write (getBound' "writeLiftIO" write)
192          (Poke $ \pf -> do x <- io; runWrite (write x) pf)
193
194-- | @writeIf p wTrue wFalse x@ creates a 'Write' with a 'Poke' equal to @wTrue
195-- x@, if @p x@ and equal to @wFalse x@ otherwise. The bound of this new
196-- 'Write' is the maximum of the bounds for either 'Write'. This yields a data
197-- independent bound, if the bound for @wTrue@ and @wFalse@ is already data
198-- independent.
199{-# INLINE writeIf #-}
200writeIf :: (a -> Bool) -> (a -> Write) -> (a -> Write) -> (a -> Write)
201writeIf p wTrue wFalse x =
202    boundedWrite (max (getBound $ wTrue x) (getBound $ wFalse x))
203                 (if p x then getPoke $ wTrue x else getPoke $ wFalse x)
204
205-- | Compare the value to a test value and use the first write action for the
206-- equal case and the second write action for the non-equal case.
207{-# INLINE writeEq #-}
208writeEq :: Eq a => a -> (a -> Write) -> (a -> Write) -> (a -> Write)
209writeEq test = writeIf (test ==)
210
211-- | TODO: Test this. It might well be too difficult to use.
212--   FIXME: Better name required!
213{-# INLINE writeOrdering #-}
214writeOrdering :: (a -> Ordering)
215              -> (a -> Write) -> (a -> Write) -> (a -> Write)
216              -> (a -> Write)
217writeOrdering ord wLT wEQ wGT x =
218    boundedWrite bound (case ord x of LT -> getPoke $ wLT x;
219                                      EQ -> getPoke $ wEQ x;
220                                      GT -> getPoke $ wGT x)
221  where
222    bound = max (getBound $ wLT x) (max (getBound $ wEQ x) (getBound $ wGT x))
223
224-- | A write combinator useful to build decision trees for deciding what value
225-- to write with a constant bound on the maximal number of bytes written.
226{-# INLINE writeOrd #-}
227writeOrd :: Ord a
228       => a
229       -> (a -> Write) -> (a -> Write) -> (a -> Write)
230       -> (a -> Write)
231writeOrd test = writeOrdering (`compare` test)
232
233-- | Create a builder that execute a single 'Write'.
234{-# INLINE fromWrite #-}
235fromWrite :: Write -> Builder
236fromWrite (Write maxSize wio) =
237    builder step
238  where
239    step k (BufferRange op ope)
240      | op `plusPtr` maxSize <= ope = do
241          op' <- runPoke wio op
242          let !br' = BufferRange op' ope
243          k br'
244      | otherwise = return $ bufferFull maxSize op (step k)
245
246{-# INLINE fromWriteSingleton #-}
247fromWriteSingleton :: (a -> Write) -> (a -> Builder)
248fromWriteSingleton write =
249    mkBuilder
250  where
251    mkBuilder x = builder step
252      where
253        step k (BufferRange op ope)
254          | op `plusPtr` maxSize <= ope = do
255              op' <- runPoke wio op
256              let !br' = BufferRange op' ope
257              k br'
258          | otherwise = return $ bufferFull maxSize op (step k)
259          where
260            Write maxSize wio = write x
261
262
263-- | Construct a 'Builder' writing a list of data one element at a time.
264fromWriteList :: (a -> Write) -> [a] -> Builder
265fromWriteList write =
266    makeBuilder
267  where
268    makeBuilder xs0 = builder $ step xs0
269      where
270        step xs1 k !(BufferRange op0 ope0) = go xs1 op0
271          where
272            go [] !op = do
273               let !br' = BufferRange op ope0
274               k br'
275
276            go xs@(x':xs') !op
277              | op `plusPtr` maxSize <= ope0 = do
278                  !op' <- runPoke wio op
279                  go xs' op'
280              | otherwise = return $ bufferFull maxSize op (step xs k)
281              where
282                Write maxSize wio = write x'
283{-# INLINE fromWriteList #-}
284
285
286
287------------------------------------------------------------------------------
288-- Writing storables
289------------------------------------------------------------------------------
290
291
292-- | Write a storable value.
293{-# INLINE writeStorable #-}
294writeStorable :: Storable a => a -> Write
295writeStorable x = exactWrite (sizeOf x) (\op -> poke (castPtr op) x)
296
297-- | A builder that serializes a storable value. No alignment is done.
298{-# INLINE fromStorable #-}
299fromStorable :: Storable a => a -> Builder
300fromStorable = fromWriteSingleton writeStorable
301
302-- | A builder that serializes a list of storable values by writing them
303-- consecutively. No alignment is done. Parsing information needs to be
304-- provided externally.
305fromStorables :: Storable a => [a] -> Builder
306fromStorables = fromWriteList writeStorable
307