1{-# LANGUAGE ScopedTypeVariables, CPP, BangPatterns, RankNTypes #-} 2#if __GLASGOW_HASKELL__ == 700 3-- This is needed as a workaround for an old bug in GHC 7.0.1 (Trac #4498) 4{-# LANGUAGE MonoPatBinds #-} 5#endif 6#if __GLASGOW_HASKELL__ >= 703 7{-# LANGUAGE Unsafe #-} 8#endif 9{-# OPTIONS_HADDOCK hide #-} 10-- | Copyright : (c) 2010 - 2011 Simon Meier 11-- License : BSD3-style (see LICENSE) 12-- 13-- Maintainer : Simon Meier <iridcode@gmail.com> 14-- Stability : unstable, private 15-- Portability : GHC 16-- 17-- *Warning:* this module is internal. If you find that you need it then please 18-- contact the maintainers and explain what you are trying to do and discuss 19-- what you would need in the public API. It is important that you do this as 20-- the module may not be exposed at all in future releases. 21-- 22-- Core types and functions for the 'Builder' monoid and its generalization, 23-- the 'Put' monad. 24-- 25-- The design of the 'Builder' monoid is optimized such that 26-- 27-- 1. buffers of arbitrary size can be filled as efficiently as possible and 28-- 29-- 2. sequencing of 'Builder's is as cheap as possible. 30-- 31-- We achieve (1) by completely handing over control over writing to the buffer 32-- to the 'BuildStep' implementing the 'Builder'. This 'BuildStep' is just told 33-- the start and the end of the buffer (represented as a 'BufferRange'). Then, 34-- the 'BuildStep' can write to as big a prefix of this 'BufferRange' in any 35-- way it desires. If the 'BuildStep' is done, the 'BufferRange' is full, or a 36-- long sequence of bytes should be inserted directly, then the 'BuildStep' 37-- signals this to its caller using a 'BuildSignal'. 38-- 39-- We achieve (2) by requiring that every 'Builder' is implemented by a 40-- 'BuildStep' that takes a continuation 'BuildStep', which it calls with the 41-- updated 'BufferRange' after it is done. Therefore, only two pointers have 42-- to be passed in a function call to implement concatenation of 'Builder's. 43-- Moreover, many 'Builder's are completely inlined, which enables the compiler 44-- to sequence them without a function call and with no boxing at all. 45-- 46-- This design gives the implementation of a 'Builder' full access to the 'IO' 47-- monad. Therefore, utmost care has to be taken to not overwrite anything 48-- outside the given 'BufferRange's. Moreover, further care has to be taken to 49-- ensure that 'Builder's and 'Put's are referentially transparent. See the 50-- comments of the 'builder' and 'put' functions for further information. 51-- Note that there are /no safety belts/ at all, when implementing a 'Builder' 52-- using an 'IO' action: you are writing code that might enable the next 53-- buffer-overflow attack on a Haskell server! 54-- 55module Data.ByteString.Builder.Internal ( 56 -- * Buffer management 57 Buffer(..) 58 , BufferRange(..) 59 , newBuffer 60 , bufferSize 61 , byteStringFromBuffer 62 63 , ChunkIOStream(..) 64 , buildStepToCIOS 65 , ciosUnitToLazyByteString 66 , ciosToLazyByteString 67 68 -- * Build signals and steps 69 , BuildSignal 70 , BuildStep 71 , finalBuildStep 72 73 , done 74 , bufferFull 75 , insertChunk 76 77 , fillWithBuildStep 78 79 -- * The Builder monoid 80 , Builder 81 , builder 82 , runBuilder 83 , runBuilderWith 84 85 -- ** Primitive combinators 86 , empty 87 , append 88 , flush 89 , ensureFree 90 -- , sizedChunksInsert 91 92 , byteStringCopy 93 , byteStringInsert 94 , byteStringThreshold 95 96 , lazyByteStringCopy 97 , lazyByteStringInsert 98 , lazyByteStringThreshold 99 100 , shortByteString 101 102 , maximalCopySize 103 , byteString 104 , lazyByteString 105 106 -- ** Execution 107 , toLazyByteStringWith 108 , AllocationStrategy 109 , safeStrategy 110 , untrimmedStrategy 111 , customStrategy 112 , L.smallChunkSize 113 , L.defaultChunkSize 114 , L.chunkOverhead 115 116 -- * The Put monad 117 , Put 118 , put 119 , runPut 120 121 -- ** Execution 122 , putToLazyByteString 123 , putToLazyByteStringWith 124 , hPut 125 126 -- ** Conversion to and from Builders 127 , putBuilder 128 , fromPut 129 130 -- -- ** Lifting IO actions 131 -- , putLiftIO 132 133) where 134 135import Control.Arrow (second) 136 137#if MIN_VERSION_base(4,9,0) 138import Data.Semigroup (Semigroup((<>))) 139#endif 140#if !(MIN_VERSION_base(4,8,0)) 141import Data.Monoid 142import Control.Applicative (Applicative(..),(<$>)) 143#endif 144 145import qualified Data.ByteString as S 146import qualified Data.ByteString.Internal as S 147import qualified Data.ByteString.Lazy.Internal as L 148import qualified Data.ByteString.Short.Internal as Sh 149 150#if __GLASGOW_HASKELL__ >= 611 151import qualified GHC.IO.Buffer as IO (Buffer(..), newByteBuffer) 152import GHC.IO.Handle.Internals (wantWritableHandle, flushWriteBuffer) 153import GHC.IO.Handle.Types (Handle__, haByteBuffer, haBufferMode) 154import System.IO (hFlush, BufferMode(..)) 155import Data.IORef 156#else 157import qualified Data.ByteString.Lazy as L 158#endif 159import System.IO (Handle) 160 161#if MIN_VERSION_base(4,4,0) 162#if MIN_VERSION_base(4,7,0) 163import Foreign 164#else 165import Foreign hiding (unsafeForeignPtrToPtr) 166#endif 167import Foreign.ForeignPtr.Unsafe (unsafeForeignPtrToPtr) 168import System.IO.Unsafe (unsafeDupablePerformIO) 169#else 170import Foreign 171import GHC.IO (unsafeDupablePerformIO) 172#endif 173 174------------------------------------------------------------------------------ 175-- Buffers 176------------------------------------------------------------------------------ 177-- | A range of bytes in a buffer represented by the pointer to the first byte 178-- of the range and the pointer to the first byte /after/ the range. 179data BufferRange = BufferRange {-# UNPACK #-} !(Ptr Word8) -- First byte of range 180 {-# UNPACK #-} !(Ptr Word8) -- First byte /after/ range 181 182-- | A 'Buffer' together with the 'BufferRange' of free bytes. The filled 183-- space starts at offset 0 and ends at the first free byte. 184data Buffer = Buffer {-# UNPACK #-} !(ForeignPtr Word8) 185 {-# UNPACK #-} !BufferRange 186 187 188-- | Combined size of the filled and free space in the buffer. 189{-# INLINE bufferSize #-} 190bufferSize :: Buffer -> Int 191bufferSize (Buffer fpbuf (BufferRange _ ope)) = 192 ope `minusPtr` unsafeForeignPtrToPtr fpbuf 193 194-- | Allocate a new buffer of the given size. 195{-# INLINE newBuffer #-} 196newBuffer :: Int -> IO Buffer 197newBuffer size = do 198 fpbuf <- S.mallocByteString size 199 let pbuf = unsafeForeignPtrToPtr fpbuf 200 return $! Buffer fpbuf (BufferRange pbuf (pbuf `plusPtr` size)) 201 202-- | Convert the filled part of a 'Buffer' to a strict 'S.ByteString'. 203{-# INLINE byteStringFromBuffer #-} 204byteStringFromBuffer :: Buffer -> S.ByteString 205byteStringFromBuffer (Buffer fpbuf (BufferRange op _)) = 206 S.PS fpbuf 0 (op `minusPtr` unsafeForeignPtrToPtr fpbuf) 207 208-- | Prepend the filled part of a 'Buffer' to a lazy 'L.ByteString' 209-- trimming it if necessary. 210{-# INLINE trimmedChunkFromBuffer #-} 211trimmedChunkFromBuffer :: AllocationStrategy -> Buffer 212 -> L.ByteString -> L.ByteString 213trimmedChunkFromBuffer (AllocationStrategy _ _ trim) buf k 214 | S.null bs = k 215 | trim (S.length bs) (bufferSize buf) = L.Chunk (S.copy bs) k 216 | otherwise = L.Chunk bs k 217 where 218 bs = byteStringFromBuffer buf 219 220------------------------------------------------------------------------------ 221-- Chunked IO Stream 222------------------------------------------------------------------------------ 223 224-- | A stream of chunks that are constructed in the 'IO' monad. 225-- 226-- This datatype serves as the common interface for the buffer-by-buffer 227-- execution of a 'BuildStep' by 'buildStepToCIOS'. Typical users of this 228-- interface are 'ciosToLazyByteString' or iteratee-style libraries like 229-- @enumerator@. 230data ChunkIOStream a = 231 Finished Buffer a 232 -- ^ The partially filled last buffer together with the result. 233 | Yield1 S.ByteString (IO (ChunkIOStream a)) 234 -- ^ Yield a /non-empty/ strict 'S.ByteString'. 235 236-- | A smart constructor for yielding one chunk that ignores the chunk if 237-- it is empty. 238{-# INLINE yield1 #-} 239yield1 :: S.ByteString -> IO (ChunkIOStream a) -> IO (ChunkIOStream a) 240yield1 bs cios | S.null bs = cios 241 | otherwise = return $ Yield1 bs cios 242 243-- | Convert a @'ChunkIOStream' ()@ to a lazy 'L.ByteString' using 244-- 'unsafeDupablePerformIO'. 245{-# INLINE ciosUnitToLazyByteString #-} 246ciosUnitToLazyByteString :: AllocationStrategy 247 -> L.ByteString -> ChunkIOStream () -> L.ByteString 248ciosUnitToLazyByteString strategy k = go 249 where 250 go (Finished buf _) = trimmedChunkFromBuffer strategy buf k 251 go (Yield1 bs io) = L.Chunk bs $ unsafeDupablePerformIO (go <$> io) 252 253-- | Convert a 'ChunkIOStream' to a lazy tuple of the result and the written 254-- 'L.ByteString' using 'unsafeDupablePerformIO'. 255{-# INLINE ciosToLazyByteString #-} 256ciosToLazyByteString :: AllocationStrategy 257 -> (a -> (b, L.ByteString)) 258 -> ChunkIOStream a 259 -> (b, L.ByteString) 260ciosToLazyByteString strategy k = 261 go 262 where 263 go (Finished buf x) = 264 second (trimmedChunkFromBuffer strategy buf) $ k x 265 go (Yield1 bs io) = second (L.Chunk bs) $ unsafeDupablePerformIO (go <$> io) 266 267------------------------------------------------------------------------------ 268-- Build signals 269------------------------------------------------------------------------------ 270 271-- | 'BuildStep's may be called *multiple times* and they must not rise an 272-- async. exception. 273type BuildStep a = BufferRange -> IO (BuildSignal a) 274 275-- | 'BuildSignal's abstract signals to the caller of a 'BuildStep'. There are 276-- three signals: 'done', 'bufferFull', or 'insertChunks signals 277data BuildSignal a = 278 Done {-# UNPACK #-} !(Ptr Word8) a 279 | BufferFull 280 {-# UNPACK #-} !Int 281 {-# UNPACK #-} !(Ptr Word8) 282 (BuildStep a) 283 | InsertChunk 284 {-# UNPACK #-} !(Ptr Word8) 285 S.ByteString 286 (BuildStep a) 287 288-- | Signal that the current 'BuildStep' is done and has computed a value. 289{-# INLINE done #-} 290done :: Ptr Word8 -- ^ Next free byte in current 'BufferRange' 291 -> a -- ^ Computed value 292 -> BuildSignal a 293done = Done 294 295-- | Signal that the current buffer is full. 296{-# INLINE bufferFull #-} 297bufferFull :: Int 298 -- ^ Minimal size of next 'BufferRange'. 299 -> Ptr Word8 300 -- ^ Next free byte in current 'BufferRange'. 301 -> BuildStep a 302 -- ^ 'BuildStep' to run on the next 'BufferRange'. This 'BuildStep' 303 -- may assume that it is called with a 'BufferRange' of at least the 304 -- required minimal size; i.e., the caller of this 'BuildStep' must 305 -- guarantee this. 306 -> BuildSignal a 307bufferFull = BufferFull 308 309 310-- | Signal that a 'S.ByteString' chunk should be inserted directly. 311{-# INLINE insertChunk #-} 312insertChunk :: Ptr Word8 313 -- ^ Next free byte in current 'BufferRange' 314 -> S.ByteString 315 -- ^ Chunk to insert. 316 -> BuildStep a 317 -- ^ 'BuildStep' to run on next 'BufferRange' 318 -> BuildSignal a 319insertChunk op bs = InsertChunk op bs 320 321 322-- | Fill a 'BufferRange' using a 'BuildStep'. 323{-# INLINE fillWithBuildStep #-} 324fillWithBuildStep 325 :: BuildStep a 326 -- ^ Build step to use for filling the 'BufferRange'. 327 -> (Ptr Word8 -> a -> IO b) 328 -- ^ Handling the 'done' signal 329 -> (Ptr Word8 -> Int -> BuildStep a -> IO b) 330 -- ^ Handling the 'bufferFull' signal 331 -> (Ptr Word8 -> S.ByteString -> BuildStep a -> IO b) 332 -- ^ Handling the 'insertChunk' signal 333 -> BufferRange 334 -- ^ Buffer range to fill. 335 -> IO b 336 -- ^ Value computed while filling this 'BufferRange'. 337fillWithBuildStep step fDone fFull fChunk !br = do 338 signal <- step br 339 case signal of 340 Done op x -> fDone op x 341 BufferFull minSize op nextStep -> fFull op minSize nextStep 342 InsertChunk op bs nextStep -> fChunk op bs nextStep 343 344 345------------------------------------------------------------------------------ 346-- The 'Builder' monoid 347------------------------------------------------------------------------------ 348 349-- | 'Builder's denote sequences of bytes. 350-- They are 'Monoid's where 351-- 'mempty' is the zero-length sequence and 352-- 'mappend' is concatenation, which runs in /O(1)/. 353newtype Builder = Builder (forall r. BuildStep r -> BuildStep r) 354 355-- | Construct a 'Builder'. In contrast to 'BuildStep's, 'Builder's are 356-- referentially transparent. 357{-# INLINE builder #-} 358builder :: (forall r. BuildStep r -> BuildStep r) 359 -- ^ A function that fills a 'BufferRange', calls the continuation with 360 -- the updated 'BufferRange' once its done, and signals its caller how 361 -- to proceed using 'done', 'bufferFull', or 'insertChunk'. 362 -- 363 -- This function must be referentially transparent; i.e., calling it 364 -- multiple times with equally sized 'BufferRange's must result in the 365 -- same sequence of bytes being written. If you need mutable state, 366 -- then you must allocate it anew upon each call of this function. 367 -- Moroever, this function must call the continuation once its done. 368 -- Otherwise, concatenation of 'Builder's does not work. Finally, this 369 -- function must write to all bytes that it claims it has written. 370 -- Otherwise, the resulting 'Builder' is not guaranteed to be 371 -- referentially transparent and sensitive data might leak. 372 -> Builder 373builder = Builder 374 375-- | The final build step that returns the 'done' signal. 376finalBuildStep :: BuildStep () 377finalBuildStep !(BufferRange op _) = return $ Done op () 378 379-- | Run a 'Builder' with the 'finalBuildStep'. 380{-# INLINE runBuilder #-} 381runBuilder :: Builder -- ^ 'Builder' to run 382 -> BuildStep () -- ^ 'BuildStep' that writes the byte stream of this 383 -- 'Builder' and signals 'done' upon completion. 384runBuilder b = runBuilderWith b finalBuildStep 385 386-- | Run a 'Builder'. 387{-# INLINE runBuilderWith #-} 388runBuilderWith :: Builder -- ^ 'Builder' to run 389 -> BuildStep a -- ^ Continuation 'BuildStep' 390 -> BuildStep a 391runBuilderWith (Builder b) = b 392 393-- | The 'Builder' denoting a zero-length sequence of bytes. This function is 394-- only exported for use in rewriting rules. Use 'mempty' otherwise. 395{-# INLINE[1] empty #-} 396empty :: Builder 397empty = Builder (\cont -> (\range -> cont range)) 398-- This eta expansion (hopefully) allows GHC to worker-wrapper the 399-- 'BufferRange' in the 'empty' base case of loops (since 400-- worker-wrapper requires (TODO: verify this) that all paths match 401-- against the wrapped argument. 402 403-- | Concatenate two 'Builder's. This function is only exported for use in rewriting 404-- rules. Use 'mappend' otherwise. 405{-# INLINE[1] append #-} 406append :: Builder -> Builder -> Builder 407append (Builder b1) (Builder b2) = Builder $ b1 . b2 408 409#if MIN_VERSION_base(4,9,0) 410instance Semigroup Builder where 411 {-# INLINE (<>) #-} 412 (<>) = append 413#endif 414 415instance Monoid Builder where 416 {-# INLINE mempty #-} 417 mempty = empty 418 {-# INLINE mappend #-} 419#if MIN_VERSION_base(4,9,0) 420 mappend = (<>) 421#else 422 mappend = append 423#endif 424 {-# INLINE mconcat #-} 425 mconcat = foldr mappend mempty 426 427-- | Flush the current buffer. This introduces a chunk boundary. 428{-# INLINE flush #-} 429flush :: Builder 430flush = builder step 431 where 432 step k !(BufferRange op _) = return $ insertChunk op S.empty k 433 434 435------------------------------------------------------------------------------ 436-- Put 437------------------------------------------------------------------------------ 438 439-- | A 'Put' action denotes a computation of a value that writes a stream of 440-- bytes as a side-effect. 'Put's are strict in their side-effect; i.e., the 441-- stream of bytes will always be written before the computed value is 442-- returned. 443-- 444-- 'Put's are a generalization of 'Builder's. The typical use case is the 445-- implementation of an encoding that might fail (e.g., an interface to the 446-- 'zlib' compression library or the conversion from Base64 encoded data to 447-- 8-bit data). For a 'Builder', the only way to handle and report such a 448-- failure is ignore it or call 'error'. In contrast, 'Put' actions are 449-- expressive enough to allow reportng and handling such a failure in a pure 450-- fashion. 451-- 452-- @'Put' ()@ actions are isomorphic to 'Builder's. The functions 'putBuilder' 453-- and 'fromPut' convert between these two types. Where possible, you should 454-- use 'Builder's, as sequencing them is slightly cheaper than sequencing 455-- 'Put's because they do not carry around a computed value. 456newtype Put a = Put { unPut :: forall r. (a -> BuildStep r) -> BuildStep r } 457 458-- | Construct a 'Put' action. In contrast to 'BuildStep's, 'Put's are 459-- referentially transparent in the sense that sequencing the same 'Put' 460-- multiple times yields every time the same value with the same side-effect. 461{-# INLINE put #-} 462put :: (forall r. (a -> BuildStep r) -> BuildStep r) 463 -- ^ A function that fills a 'BufferRange', calls the continuation with 464 -- the updated 'BufferRange' and its computed value once its done, and 465 -- signals its caller how to proceed using 'done', 'bufferFull', or 466 -- 'insertChunk' signals. 467 -- 468 -- This function must be referentially transparent; i.e., calling it 469 -- multiple times with equally sized 'BufferRange's must result in the 470 -- same sequence of bytes being written and the same value being 471 -- computed. If you need mutable state, then you must allocate it anew 472 -- upon each call of this function. Moroever, this function must call 473 -- the continuation once its done. Otherwise, monadic sequencing of 474 -- 'Put's does not work. Finally, this function must write to all bytes 475 -- that it claims it has written. Otherwise, the resulting 'Put' is 476 -- not guaranteed to be referentially transparent and sensitive data 477 -- might leak. 478 -> Put a 479put = Put 480 481-- | Run a 'Put'. 482{-# INLINE runPut #-} 483runPut :: Put a -- ^ Put to run 484 -> BuildStep a -- ^ 'BuildStep' that first writes the byte stream of 485 -- this 'Put' and then yields the computed value using 486 -- the 'done' signal. 487runPut (Put p) = p $ \x (BufferRange op _) -> return $ Done op x 488 489instance Functor Put where 490 fmap f p = Put $ \k -> unPut p (\x -> k (f x)) 491 {-# INLINE fmap #-} 492 493-- | Synonym for '<*' from 'Applicative'; used in rewriting rules. 494{-# INLINE[1] ap_l #-} 495ap_l :: Put a -> Put b -> Put a 496ap_l (Put a) (Put b) = Put $ \k -> a (\a' -> b (\_ -> k a')) 497 498-- | Synonym for '*>' from 'Applicative' and '>>' from 'Monad'; used in 499-- rewriting rules. 500{-# INLINE[1] ap_r #-} 501ap_r :: Put a -> Put b -> Put b 502ap_r (Put a) (Put b) = Put $ \k -> a (\_ -> b k) 503 504instance Applicative Put where 505 {-# INLINE pure #-} 506 pure x = Put $ \k -> k x 507 {-# INLINE (<*>) #-} 508 Put f <*> Put a = Put $ \k -> f (\f' -> a (\a' -> k (f' a'))) 509 {-# INLINE (<*) #-} 510 (<*) = ap_l 511 {-# INLINE (*>) #-} 512 (*>) = ap_r 513 514instance Monad Put where 515 {-# INLINE return #-} 516 return = pure 517 {-# INLINE (>>=) #-} 518 Put m >>= f = Put $ \k -> m (\m' -> unPut (f m') k) 519 {-# INLINE (>>) #-} 520 (>>) = (*>) 521 522-- Conversion between Put and Builder 523------------------------------------- 524 525-- | Run a 'Builder' as a side-effect of a @'Put' ()@ action. 526{-# INLINE[1] putBuilder #-} 527putBuilder :: Builder -> Put () 528putBuilder (Builder b) = Put $ \k -> b (k ()) 529 530-- | Convert a @'Put' ()@ action to a 'Builder'. 531{-# INLINE fromPut #-} 532fromPut :: Put () -> Builder 533fromPut (Put p) = Builder $ \k -> p (\_ -> k) 534 535-- We rewrite consecutive uses of 'putBuilder' such that the append of the 536-- involved 'Builder's is used. This can significantly improve performance, 537-- when the bound-checks of the concatenated builders are fused. 538 539-- ap_l rules 540{-# RULES 541 542"ap_l/putBuilder" forall b1 b2. 543 ap_l (putBuilder b1) (putBuilder b2) 544 = putBuilder (append b1 b2) 545 546"ap_l/putBuilder/assoc_r" forall b1 b2 (p :: Put a). 547 ap_l (putBuilder b1) (ap_l (putBuilder b2) p) 548 = ap_l (putBuilder (append b1 b2)) p 549 550"ap_l/putBuilder/assoc_l" forall (p :: Put a) b1 b2. 551 ap_l (ap_l p (putBuilder b1)) (putBuilder b2) 552 = ap_l p (putBuilder (append b1 b2)) 553 #-} 554 555-- ap_r rules 556{-# RULES 557 558"ap_r/putBuilder" forall b1 b2. 559 ap_r (putBuilder b1) (putBuilder b2) 560 = putBuilder (append b1 b2) 561 562"ap_r/putBuilder/assoc_r" forall b1 b2 (p :: Put a). 563 ap_r (putBuilder b1) (ap_r (putBuilder b2) p) 564 = ap_r (putBuilder (append b1 b2)) p 565 566"ap_r/putBuilder/assoc_l" forall (p :: Put a) b1 b2. 567 ap_r (ap_r p (putBuilder b1)) (putBuilder b2) 568 = ap_r p (putBuilder (append b1 b2)) 569 570 #-} 571 572-- combined ap_l/ap_r rules 573{-# RULES 574 575"ap_l/ap_r/putBuilder/assoc_r" forall b1 b2 (p :: Put a). 576 ap_l (putBuilder b1) (ap_r (putBuilder b2) p) 577 = ap_l (putBuilder (append b1 b2)) p 578 579"ap_r/ap_l/putBuilder/assoc_r" forall b1 b2 (p :: Put a). 580 ap_r (putBuilder b1) (ap_l (putBuilder b2) p) 581 = ap_l (putBuilder (append b1 b2)) p 582 583"ap_l/ap_r/putBuilder/assoc_l" forall (p :: Put a) b1 b2. 584 ap_l (ap_r p (putBuilder b1)) (putBuilder b2) 585 = ap_r p (putBuilder (append b1 b2)) 586 587"ap_r/ap_l/putBuilder/assoc_l" forall (p :: Put a) b1 b2. 588 ap_r (ap_l p (putBuilder b1)) (putBuilder b2) 589 = ap_r p (putBuilder (append b1 b2)) 590 591 #-} 592 593 594-- Lifting IO actions 595--------------------- 596 597{- 598-- | Lift an 'IO' action to a 'Put' action. 599{-# INLINE putLiftIO #-} 600putLiftIO :: IO a -> Put a 601putLiftIO io = put $ \k br -> io >>= (`k` br) 602-} 603 604 605------------------------------------------------------------------------------ 606-- Executing a Put directly on a buffered Handle 607------------------------------------------------------------------------------ 608 609-- | Run a 'Put' action redirecting the produced output to a 'Handle'. 610-- 611-- The output is buffered using the 'Handle's associated buffer. If this 612-- buffer is too small to execute one step of the 'Put' action, then 613-- it is replaced with a large enough buffer. 614hPut :: forall a. Handle -> Put a -> IO a 615#if __GLASGOW_HASKELL__ >= 611 616hPut h p = do 617 fillHandle 1 (runPut p) 618 where 619 fillHandle :: Int -> BuildStep a -> IO a 620 fillHandle !minFree step = do 621 next <- wantWritableHandle "hPut" h fillHandle_ 622 next 623 where 624 -- | We need to return an inner IO action that is executed outside 625 -- the lock taken on the Handle for two reasons: 626 -- 627 -- 1. GHC.IO.Handle.Internals mentions in "Note [async]" that 628 -- we should never do any side-effecting operations before 629 -- an interuptible operation that may raise an async. exception 630 -- as long as we are inside 'wantWritableHandle' and the like. 631 -- We possibly run the interuptible 'flushWriteBuffer' right at 632 -- the start of 'fillHandle', hence entering it a second time is 633 -- not safe, as it could lead to a 'BuildStep' being run twice. 634 -- 635 -- FIXME (SM): Adapt this function or at least its documentation, 636 -- as it is OK to run a 'BuildStep' twice. We dropped this 637 -- requirement in favor of being able to use 638 -- 'unsafeDupablePerformIO' and the speed improvement that it 639 -- brings. 640 -- 641 -- 2. We use the 'S.hPut' function to also write to the handle. 642 -- This function tries to take the same lock taken by 643 -- 'wantWritableHandle'. Therefore, we cannot call 'S.hPut' 644 -- inside 'wantWritableHandle'. 645 -- 646 fillHandle_ :: Handle__ -> IO (IO a) 647 fillHandle_ h_ = do 648 makeSpace =<< readIORef refBuf 649 fillBuffer =<< readIORef refBuf 650 where 651 refBuf = haByteBuffer h_ 652 freeSpace buf = IO.bufSize buf - IO.bufR buf 653 654 makeSpace buf 655 | IO.bufSize buf < minFree = do 656 flushWriteBuffer h_ 657 s <- IO.bufState <$> readIORef refBuf 658 IO.newByteBuffer minFree s >>= writeIORef refBuf 659 660 | freeSpace buf < minFree = flushWriteBuffer h_ 661 | otherwise = 662#if __GLASGOW_HASKELL__ >= 613 663 return () 664#else 665 -- required for ghc-6.12 666 flushWriteBuffer h_ 667#endif 668 669 fillBuffer buf 670 | freeSpace buf < minFree = 671 error $ unlines 672 [ "Data.ByteString.Builder.Internal.hPut: internal error." 673 , " Not enough space after flush." 674 , " required: " ++ show minFree 675 , " free: " ++ show (freeSpace buf) 676 ] 677 | otherwise = do 678 let !br = BufferRange op (pBuf `plusPtr` IO.bufSize buf) 679 res <- fillWithBuildStep step doneH fullH insertChunkH br 680 touchForeignPtr fpBuf 681 return res 682 where 683 fpBuf = IO.bufRaw buf 684 pBuf = unsafeForeignPtrToPtr fpBuf 685 op = pBuf `plusPtr` IO.bufR buf 686 687 {-# INLINE updateBufR #-} 688 updateBufR op' = do 689 let !off' = op' `minusPtr` pBuf 690 !buf' = buf {IO.bufR = off'} 691 writeIORef refBuf buf' 692 693 doneH op' x = do 694 updateBufR op' 695 -- We must flush if this Handle is set to NoBuffering. 696 -- If it is set to LineBuffering, be conservative and 697 -- flush anyway (we didn't check for newlines in the data). 698 -- Flushing must happen outside this 'wantWriteableHandle' 699 -- due to the possible async. exception. 700 case haBufferMode h_ of 701 BlockBuffering _ -> return $ return x 702 _line_or_no_buffering -> return $ hFlush h >> return x 703 704 fullH op' minSize nextStep = do 705 updateBufR op' 706 return $ fillHandle minSize nextStep 707 -- 'fillHandle' will flush the buffer (provided there is 708 -- really less than 'minSize' space left) before executing 709 -- the 'nextStep'. 710 711 insertChunkH op' bs nextStep = do 712 updateBufR op' 713 return $ do 714 S.hPut h bs 715 fillHandle 1 nextStep 716#else 717hPut h p = 718 go =<< buildStepToCIOS strategy (runPut p) 719 where 720 strategy = untrimmedStrategy L.smallChunkSize L.defaultChunkSize 721 722 go (Finished buf x) = S.hPut h (byteStringFromBuffer buf) >> return x 723 go (Yield1 bs io) = S.hPut h bs >> io >>= go 724#endif 725 726-- | Execute a 'Put' and return the computed result and the bytes 727-- written during the computation as a lazy 'L.ByteString'. 728-- 729-- This function is strict in the computed result and lazy in the writing of 730-- the bytes. For example, given 731-- 732-- @ 733--infinitePut = sequence_ (repeat (putBuilder (word8 1))) >> return 0 734-- @ 735-- 736-- evaluating the expression 737-- 738-- @ 739--fst $ putToLazyByteString infinitePut 740-- @ 741-- 742-- does not terminate, while evaluating the expression 743-- 744-- @ 745--L.head $ snd $ putToLazyByteString infinitePut 746-- @ 747-- 748-- does terminate and yields the value @1 :: Word8@. 749-- 750-- An illustrative example for these strictness properties is the 751-- implementation of Base64 decoding (<http://en.wikipedia.org/wiki/Base64>). 752-- 753-- @ 754--type DecodingState = ... 755-- 756--decodeBase64 :: 'S.ByteString' -> DecodingState -> 'Put' (Maybe DecodingState) 757--decodeBase64 = ... 758-- @ 759-- 760-- The above function takes a strict 'S.ByteString' supposed to represent 761-- Base64 encoded data and the current decoding state. 762-- It writes the decoded bytes as the side-effect of the 'Put' and returns the 763-- new decoding state, if the decoding of all data in the 'S.ByteString' was 764-- successful. The checking if the strict 'S.ByteString' represents Base64 765-- encoded data and the actual decoding are fused. This makes the common case, 766-- where all data represents Base64 encoded data, more efficient. It also 767-- implies that all data must be decoded before the final decoding 768-- state can be returned. 'Put's are intended for implementing such fused 769-- checking and decoding/encoding, which is reflected in their strictness 770-- properties. 771{-# NOINLINE putToLazyByteString #-} 772putToLazyByteString 773 :: Put a -- ^ 'Put' to execute 774 -> (a, L.ByteString) -- ^ Result and lazy 'L.ByteString' 775 -- written as its side-effect 776putToLazyByteString = putToLazyByteStringWith 777 (safeStrategy L.smallChunkSize L.defaultChunkSize) (\x -> (x, L.Empty)) 778 779 780-- | Execute a 'Put' with a buffer-allocation strategy and a continuation. For 781-- example, 'putToLazyByteString' is implemented as follows. 782-- 783-- @ 784--putToLazyByteString = 'putToLazyByteStringWith' 785-- ('safeStrategy' 'L.smallChunkSize' 'L.defaultChunkSize') (\x -> (x, L.empty)) 786-- @ 787-- 788{-# INLINE putToLazyByteStringWith #-} 789putToLazyByteStringWith 790 :: AllocationStrategy 791 -- ^ Buffer allocation strategy to use 792 -> (a -> (b, L.ByteString)) 793 -- ^ Continuation to use for computing the final result and the tail of 794 -- its side-effect (the written bytes). 795 -> Put a 796 -- ^ 'Put' to execute 797 -> (b, L.ByteString) 798 -- ^ Resulting lazy 'L.ByteString' 799putToLazyByteStringWith strategy k p = 800 ciosToLazyByteString strategy k $ unsafeDupablePerformIO $ 801 buildStepToCIOS strategy (runPut p) 802 803 804 805------------------------------------------------------------------------------ 806-- ByteString insertion / controlling chunk boundaries 807------------------------------------------------------------------------------ 808 809-- Raw memory 810------------- 811 812-- | Ensure that there are at least 'n' free bytes for the following 'Builder'. 813{-# INLINE ensureFree #-} 814ensureFree :: Int -> Builder 815ensureFree minFree = 816 builder step 817 where 818 step k br@(BufferRange op ope) 819 | ope `minusPtr` op < minFree = return $ bufferFull minFree op k 820 | otherwise = k br 821 822-- | Copy the bytes from a 'BufferRange' into the output stream. 823wrappedBytesCopyStep :: BufferRange -- ^ Input 'BufferRange'. 824 -> BuildStep a -> BuildStep a 825wrappedBytesCopyStep !(BufferRange ip0 ipe) k = 826 go ip0 827 where 828 go !ip !(BufferRange op ope) 829 | inpRemaining <= outRemaining = do 830 copyBytes op ip inpRemaining 831 let !br' = BufferRange (op `plusPtr` inpRemaining) ope 832 k br' 833 | otherwise = do 834 copyBytes op ip outRemaining 835 let !ip' = ip `plusPtr` outRemaining 836 return $ bufferFull 1 ope (go ip') 837 where 838 outRemaining = ope `minusPtr` op 839 inpRemaining = ipe `minusPtr` ip 840 841 842-- Strict ByteStrings 843------------------------------------------------------------------------------ 844 845 846-- | Construct a 'Builder' that copies the strict 'S.ByteString's, if it is 847-- smaller than the treshold, and inserts it directly otherwise. 848-- 849-- For example, @byteStringThreshold 1024@ copies strict 'S.ByteString's whose size 850-- is less or equal to 1kb, and inserts them directly otherwise. This implies 851-- that the average chunk-size of the generated lazy 'L.ByteString' may be as 852-- low as 513 bytes, as there could always be just a single byte between the 853-- directly inserted 1025 byte, strict 'S.ByteString's. 854-- 855{-# INLINE byteStringThreshold #-} 856byteStringThreshold :: Int -> S.ByteString -> Builder 857byteStringThreshold maxCopySize = 858 \bs -> builder $ step bs 859 where 860 step !bs@(S.PS _ _ len) !k br@(BufferRange !op _) 861 | len <= maxCopySize = byteStringCopyStep bs k br 862 | otherwise = return $ insertChunk op bs k 863 864-- | Construct a 'Builder' that copies the strict 'S.ByteString'. 865-- 866-- Use this function to create 'Builder's from smallish (@<= 4kb@) 867-- 'S.ByteString's or if you need to guarantee that the 'S.ByteString' is not 868-- shared with the chunks generated by the 'Builder'. 869-- 870{-# INLINE byteStringCopy #-} 871byteStringCopy :: S.ByteString -> Builder 872byteStringCopy = \bs -> builder $ byteStringCopyStep bs 873 874{-# INLINE byteStringCopyStep #-} 875byteStringCopyStep :: S.ByteString -> BuildStep a -> BuildStep a 876byteStringCopyStep (S.PS ifp ioff isize) !k0 br0@(BufferRange op ope) 877 -- Ensure that the common case is not recursive and therefore yields 878 -- better code. 879 | op' <= ope = do copyBytes op ip isize 880 touchForeignPtr ifp 881 k0 (BufferRange op' ope) 882 | otherwise = do wrappedBytesCopyStep (BufferRange ip ipe) k br0 883 where 884 op' = op `plusPtr` isize 885 ip = unsafeForeignPtrToPtr ifp `plusPtr` ioff 886 ipe = ip `plusPtr` isize 887 k br = do touchForeignPtr ifp -- input consumed: OK to release here 888 k0 br 889 890-- | Construct a 'Builder' that always inserts the strict 'S.ByteString' 891-- directly as a chunk. 892-- 893-- This implies flushing the output buffer, even if it contains just 894-- a single byte. You should therefore use 'byteStringInsert' only for large 895-- (@> 8kb@) 'S.ByteString's. Otherwise, the generated chunks are too 896-- fragmented to be processed efficiently afterwards. 897-- 898{-# INLINE byteStringInsert #-} 899byteStringInsert :: S.ByteString -> Builder 900byteStringInsert = 901 \bs -> builder $ \k (BufferRange op _) -> return $ insertChunk op bs k 902 903-- Short bytestrings 904------------------------------------------------------------------------------ 905 906-- | Construct a 'Builder' that copies the 'SH.ShortByteString'. 907-- 908{-# INLINE shortByteString #-} 909shortByteString :: Sh.ShortByteString -> Builder 910shortByteString = \sbs -> builder $ shortByteStringCopyStep sbs 911 912-- | Copy the bytes from a 'SH.ShortByteString' into the output stream. 913{-# INLINE shortByteStringCopyStep #-} 914shortByteStringCopyStep :: Sh.ShortByteString -- ^ Input 'SH.ShortByteString'. 915 -> BuildStep a -> BuildStep a 916shortByteStringCopyStep !sbs k = 917 go 0 (Sh.length sbs) 918 where 919 go !ip !ipe !(BufferRange op ope) 920 | inpRemaining <= outRemaining = do 921 Sh.copyToPtr sbs ip op inpRemaining 922 let !br' = BufferRange (op `plusPtr` inpRemaining) ope 923 k br' 924 | otherwise = do 925 Sh.copyToPtr sbs ip op outRemaining 926 let !ip' = ip + outRemaining 927 return $ bufferFull 1 ope (go ip' ipe) 928 where 929 outRemaining = ope `minusPtr` op 930 inpRemaining = ipe - ip 931 932 933-- Lazy bytestrings 934------------------------------------------------------------------------------ 935 936-- | Construct a 'Builder' that uses the thresholding strategy of 'byteStringThreshold' 937-- for each chunk of the lazy 'L.ByteString'. 938-- 939{-# INLINE lazyByteStringThreshold #-} 940lazyByteStringThreshold :: Int -> L.ByteString -> Builder 941lazyByteStringThreshold maxCopySize = 942 L.foldrChunks (\bs b -> byteStringThreshold maxCopySize bs `mappend` b) mempty 943 -- TODO: We could do better here. Currently, Large, Small, Large, leads to 944 -- an unnecessary copy of the 'Small' chunk. 945 946-- | Construct a 'Builder' that copies the lazy 'L.ByteString'. 947-- 948{-# INLINE lazyByteStringCopy #-} 949lazyByteStringCopy :: L.ByteString -> Builder 950lazyByteStringCopy = 951 L.foldrChunks (\bs b -> byteStringCopy bs `mappend` b) mempty 952 953-- | Construct a 'Builder' that inserts all chunks of the lazy 'L.ByteString' 954-- directly. 955-- 956{-# INLINE lazyByteStringInsert #-} 957lazyByteStringInsert :: L.ByteString -> Builder 958lazyByteStringInsert = 959 L.foldrChunks (\bs b -> byteStringInsert bs `mappend` b) mempty 960 961-- | Create a 'Builder' denoting the same sequence of bytes as a strict 962-- 'S.ByteString'. 963-- The 'Builder' inserts large 'S.ByteString's directly, but copies small ones 964-- to ensure that the generated chunks are large on average. 965-- 966{-# INLINE byteString #-} 967byteString :: S.ByteString -> Builder 968byteString = byteStringThreshold maximalCopySize 969 970-- | Create a 'Builder' denoting the same sequence of bytes as a lazy 971-- 'S.ByteString'. 972-- The 'Builder' inserts large chunks of the lazy 'L.ByteString' directly, 973-- but copies small ones to ensure that the generated chunks are large on 974-- average. 975-- 976{-# INLINE lazyByteString #-} 977lazyByteString :: L.ByteString -> Builder 978lazyByteString = lazyByteStringThreshold maximalCopySize 979-- FIXME: also insert the small chunk for [large,small,large] directly. 980-- Perhaps it makes even sense to concatenate the small chunks in 981-- [large,small,small,small,large] and insert them directly afterwards to avoid 982-- unnecessary buffer spilling. Hmm, but that uncontrollably increases latency 983-- => no good! 984 985-- | The maximal size of a 'S.ByteString' that is copied. 986-- @2 * 'L.smallChunkSize'@ to guarantee that on average a chunk is of 987-- 'L.smallChunkSize'. 988maximalCopySize :: Int 989maximalCopySize = 2 * L.smallChunkSize 990 991------------------------------------------------------------------------------ 992-- Builder execution 993------------------------------------------------------------------------------ 994 995-- | A buffer allocation strategy for executing 'Builder's. 996 997-- The strategy 998-- 999-- > 'AllocationStrategy' firstBufSize bufSize trim 1000-- 1001-- states that the first buffer is of size @firstBufSize@, all following buffers 1002-- are of size @bufSize@, and a buffer of size @n@ filled with @k@ bytes should 1003-- be trimmed iff @trim k n@ is 'True'. 1004data AllocationStrategy = AllocationStrategy 1005 (Maybe (Buffer, Int) -> IO Buffer) 1006 {-# UNPACK #-} !Int 1007 (Int -> Int -> Bool) 1008 1009-- | Create a custom allocation strategy. See the code for 'safeStrategy' and 1010-- 'untrimmedStrategy' for examples. 1011{-# INLINE customStrategy #-} 1012customStrategy 1013 :: (Maybe (Buffer, Int) -> IO Buffer) 1014 -- ^ Buffer allocation function. If 'Nothing' is given, then a new first 1015 -- buffer should be allocated. If @'Just' (oldBuf, minSize)@ is given, 1016 -- then a buffer with minimal size 'minSize' must be returned. The 1017 -- strategy may reuse the 'oldBuffer', if it can guarantee that this 1018 -- referentially transparent and 'oldBuffer' is large enough. 1019 -> Int 1020 -- ^ Default buffer size. 1021 -> (Int -> Int -> Bool) 1022 -- ^ A predicate @trim used allocated@ returning 'True', if the buffer 1023 -- should be trimmed before it is returned. 1024 -> AllocationStrategy 1025customStrategy = AllocationStrategy 1026 1027-- | Sanitize a buffer size; i.e., make it at least the size of an 'Int'. 1028{-# INLINE sanitize #-} 1029sanitize :: Int -> Int 1030sanitize = max (sizeOf (undefined :: Int)) 1031 1032-- | Use this strategy for generating lazy 'L.ByteString's whose chunks are 1033-- discarded right after they are generated. For example, if you just generate 1034-- them to write them to a network socket. 1035{-# INLINE untrimmedStrategy #-} 1036untrimmedStrategy :: Int -- ^ Size of the first buffer 1037 -> Int -- ^ Size of successive buffers 1038 -> AllocationStrategy 1039 -- ^ An allocation strategy that does not trim any of the 1040 -- filled buffers before converting it to a chunk 1041untrimmedStrategy firstSize bufSize = 1042 AllocationStrategy nextBuffer (sanitize bufSize) (\_ _ -> False) 1043 where 1044 {-# INLINE nextBuffer #-} 1045 nextBuffer Nothing = newBuffer $ sanitize firstSize 1046 nextBuffer (Just (_, minSize)) = newBuffer minSize 1047 1048 1049-- | Use this strategy for generating lazy 'L.ByteString's whose chunks are 1050-- likely to survive one garbage collection. This strategy trims buffers 1051-- that are filled less than half in order to avoid spilling too much memory. 1052{-# INLINE safeStrategy #-} 1053safeStrategy :: Int -- ^ Size of first buffer 1054 -> Int -- ^ Size of successive buffers 1055 -> AllocationStrategy 1056 -- ^ An allocation strategy that guarantees that at least half 1057 -- of the allocated memory is used for live data 1058safeStrategy firstSize bufSize = 1059 AllocationStrategy nextBuffer (sanitize bufSize) trim 1060 where 1061 trim used size = 2 * used < size 1062 {-# INLINE nextBuffer #-} 1063 nextBuffer Nothing = newBuffer $ sanitize firstSize 1064 nextBuffer (Just (_, minSize)) = newBuffer minSize 1065 1066-- | /Heavy inlining./ Execute a 'Builder' with custom execution parameters. 1067-- 1068-- This function is inlined despite its heavy code-size to allow fusing with 1069-- the allocation strategy. For example, the default 'Builder' execution 1070-- function 'toLazyByteString' is defined as follows. 1071-- 1072-- @ 1073-- {-\# NOINLINE toLazyByteString \#-} 1074-- toLazyByteString = 1075-- toLazyByteStringWith ('safeStrategy' 'L.smallChunkSize' 'L.defaultChunkSize') L.empty 1076-- @ 1077-- 1078-- where @L.empty@ is the zero-length lazy 'L.ByteString'. 1079-- 1080-- In most cases, the parameters used by 'toLazyByteString' give good 1081-- performance. A sub-performing case of 'toLazyByteString' is executing short 1082-- (<128 bytes) 'Builder's. In this case, the allocation overhead for the first 1083-- 4kb buffer and the trimming cost dominate the cost of executing the 1084-- 'Builder'. You can avoid this problem using 1085-- 1086-- >toLazyByteStringWith (safeStrategy 128 smallChunkSize) L.empty 1087-- 1088-- This reduces the allocation and trimming overhead, as all generated 1089-- 'L.ByteString's fit into the first buffer and there is no trimming 1090-- required, if more than 64 bytes and less than 128 bytes are written. 1091-- 1092{-# INLINE toLazyByteStringWith #-} 1093toLazyByteStringWith 1094 :: AllocationStrategy 1095 -- ^ Buffer allocation strategy to use 1096 -> L.ByteString 1097 -- ^ Lazy 'L.ByteString' to use as the tail of the generated lazy 1098 -- 'L.ByteString' 1099 -> Builder 1100 -- ^ 'Builder' to execute 1101 -> L.ByteString 1102 -- ^ Resulting lazy 'L.ByteString' 1103toLazyByteStringWith strategy k b = 1104 ciosUnitToLazyByteString strategy k $ unsafeDupablePerformIO $ 1105 buildStepToCIOS strategy (runBuilder b) 1106 1107-- | Convert a 'BuildStep' to a 'ChunkIOStream' stream by executing it on 1108-- 'Buffer's allocated according to the given 'AllocationStrategy'. 1109{-# INLINE buildStepToCIOS #-} 1110buildStepToCIOS 1111 :: AllocationStrategy -- ^ Buffer allocation strategy to use 1112 -> BuildStep a -- ^ 'BuildStep' to execute 1113 -> IO (ChunkIOStream a) 1114buildStepToCIOS !(AllocationStrategy nextBuffer bufSize trim) = 1115 \step -> nextBuffer Nothing >>= fill step 1116 where 1117 fill !step !buf@(Buffer fpbuf br@(BufferRange _ pe)) = do 1118 res <- fillWithBuildStep step doneH fullH insertChunkH br 1119 touchForeignPtr fpbuf 1120 return res 1121 where 1122 pbuf = unsafeForeignPtrToPtr fpbuf 1123 1124 doneH op' x = return $ 1125 Finished (Buffer fpbuf (BufferRange op' pe)) x 1126 1127 fullH op' minSize nextStep = 1128 wrapChunk op' $ const $ 1129 nextBuffer (Just (buf, max minSize bufSize)) >>= fill nextStep 1130 1131 insertChunkH op' bs nextStep = 1132 wrapChunk op' $ \isEmpty -> yield1 bs $ 1133 -- Checking for empty case avoids allocating 'n-1' empty 1134 -- buffers for 'n' insertChunkH right after each other. 1135 if isEmpty 1136 then fill nextStep buf 1137 else do buf' <- nextBuffer (Just (buf, bufSize)) 1138 fill nextStep buf' 1139 1140 -- Wrap and yield a chunk, trimming it if necesary 1141 {-# INLINE wrapChunk #-} 1142 wrapChunk !op' mkCIOS 1143 | chunkSize == 0 = mkCIOS True 1144 | trim chunkSize size = do 1145 bs <- S.create chunkSize $ \pbuf' -> 1146 copyBytes pbuf' pbuf chunkSize 1147 -- FIXME: We could reuse the trimmed buffer here. 1148 return $ Yield1 bs (mkCIOS False) 1149 | otherwise = 1150 return $ Yield1 (S.PS fpbuf 0 chunkSize) (mkCIOS False) 1151 where 1152 chunkSize = op' `minusPtr` pbuf 1153 size = pe `minusPtr` pbuf 1154