Blaze.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE RankNTypes #-}
-- | Convert a stream of blaze-builder @Builder@s into a stream of @ByteString@s.
--
-- Adapted from blaze-builder-enumerator, written by myself and Simon Meier.
--
-- Note that the functions here can work in any monad built on top of @IO@ or
-- @ST@.
module Data.Conduit.Blaze
    (

  -- * Conduits from builders to bytestrings
    builderToByteString
  , unsafeBuilderToByteString
  , builderToByteStringWith

  -- ** Flush
  , builderToByteStringFlush
  , builderToByteStringWithFlush

  -- * Buffers
  , Buffer

  -- ** Status information
  , freeSize
  , sliceSize
  , bufferSize

  -- ** Creation and modification
  , allocBuffer
  , reuseBuffer
  , nextSlice

  -- ** Conversion to bytestings
  , unsafeFreezeBuffer
  , unsafeFreezeNonEmptyBuffer

  -- * Buffer allocation strategies
  , BufferAllocStrategy
  , allNewBuffersStrategy
  , reuseBufferStrategy
    ) where

import Data.Conduit
import Control.Monad (unless, liftM)
import Control.Monad.Trans.Class (lift, MonadTrans)

import qualified Data.ByteString                   as S

import Blaze.ByteString.Builder.Internal
import Blaze.ByteString.Builder.Internal.Types
import Blaze.ByteString.Builder.Internal.Buffer
import Control.Monad.Primitive (PrimMonad, unsafePrimToPrim)
import Control.Monad.Base (MonadBase, liftBase)

unsafeLiftIO :: (MonadBase base m, PrimMonad base) => IO a -> m a
unsafeLiftIO = liftBase . unsafePrimToPrim

-- | Incrementally execute builders and pass on the filled chunks as
-- bytestrings.
builderToByteString :: (MonadBase base m, PrimMonad base) => Conduit Builder m S.ByteString
builderToByteString =
  builderToByteStringWith (allNewBuffersStrategy defaultBufferSize)

-- |
--
-- Since 0.0.2
builderToByteStringFlush :: (MonadBase base m, PrimMonad base) => Conduit (Flush Builder) m (Flush S.ByteString)
builderToByteStringFlush =
  builderToByteStringWithFlush (allNewBuffersStrategy defaultBufferSize)

-- | Incrementally execute builders on the given buffer and pass on the filled
-- chunks as bytestrings. Note that, if the given buffer is too small for the
-- execution of a build step, a larger one will be allocated.
--
-- WARNING: This conduit yields bytestrings that are NOT
-- referentially transparent. Their content will be overwritten as soon
-- as control is returned from the inner sink!
unsafeBuilderToByteString :: (MonadBase base m, PrimMonad base)
                          => IO Buffer  -- action yielding the inital buffer.
                          -> Conduit Builder m S.ByteString
unsafeBuilderToByteString = builderToByteStringWith . reuseBufferStrategy


-- | A conduit that incrementally executes builders and passes on the
-- filled chunks as bytestrings to an inner sink.
--
-- INV: All bytestrings passed to the inner sink are non-empty.
builderToByteStringWith :: (MonadBase base m, PrimMonad base)
                        => BufferAllocStrategy
                        -> Conduit Builder m S.ByteString
builderToByteStringWith =
    helper (liftM (fmap Chunk) await) yield'
  where
    yield' Flush = return ()
    yield' (Chunk bs) = yield bs

-- |
--
-- Since 0.0.2
builderToByteStringWithFlush
    :: (MonadBase base m, PrimMonad base)
    => BufferAllocStrategy
    -> Conduit (Flush Builder) m (Flush S.ByteString)
builderToByteStringWithFlush = helper await yield

helper :: (MonadBase base m, PrimMonad base, Monad (t m), MonadTrans t)
       => t m (Maybe (Flush Builder))
       -> (Flush S.ByteString -> t m ())
       -> BufferAllocStrategy
       -> t m ()
helper await' yield' (ioBufInit, nextBuf) =
    loop ioBufInit
  where
    loop ioBuf = do
        await' >>= maybe (close ioBuf) (cont' ioBuf)

    cont' ioBuf Flush = push ioBuf flush $ \ioBuf' -> yield' Flush >> loop ioBuf'
    cont' ioBuf (Chunk builder) = push ioBuf builder loop

    close ioBuf = do
        buf <- lift $ unsafeLiftIO $ ioBuf
        maybe (return ()) (yield' . Chunk) (unsafeFreezeNonEmptyBuffer buf)

    push ioBuf0 x continue = do
        go (unBuilder x (buildStep finalStep)) ioBuf0
      where
        finalStep !(BufRange pf _) = return $ Done pf ()

        go bStep ioBuf = do
            !buf   <- lift $ unsafeLiftIO $ ioBuf
            signal <- lift $ unsafeLiftIO $ execBuildStep bStep buf
            case signal of
                Done op' _ -> continue $ return $ updateEndOfSlice buf op'
                BufferFull minSize op' bStep' -> do
                    let buf' = updateEndOfSlice buf op'
                        {-# INLINE cont #-}
                        cont = 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' <- lift $ unsafeLiftIO $ nextBuf minSize buf'
                            go bStep' ioBuf'
                    case unsafeFreezeNonEmptyBuffer buf' of
                        Nothing -> return ()
                        Just bs -> yield' (Chunk bs)
                    cont
                InsertByteString op' bs bStep' -> do
                    let buf' = updateEndOfSlice buf op'
                    case unsafeFreezeNonEmptyBuffer buf' of
                        Nothing -> return ()
                        Just bs' -> yield' $ Chunk bs'
                    unless (S.null bs) $ yield' $ Chunk bs
                    lift (unsafeLiftIO $ nextBuf 1 buf') >>= go bStep'