{-# LANGUAGE CPP, BangPatterns, Rank2Types, MagicHash #-}
-- |
-- Module      : UnboxedAppend
-- Copyright   : (c) 2010 Simon Meier
-- License     : BSD3-style (see LICENSE)
-- 
-- Maintainer  : Simon Meier <iridcode@gmail.com>
-- Stability   : experimental
-- Portability : tested on GHC only
--
-- Try using unboxed pointers for the continuation calls to make abstract
-- appends go faster.
--
module UnboxedAppend where

import Foreign
import Foreign.UPtr
import Data.Monoid
import qualified Data.ByteString      as S
import qualified Data.ByteString.Lazy as L

#ifdef BYTESTRING_IN_BASE
import Data.ByteString.Base (inlinePerformIO)
import qualified Data.ByteString.Base as S
import qualified Data.ByteString.Lazy.Base as L -- FIXME: is this the right module for access to 'Chunks'?
#else
import Data.ByteString.Internal (inlinePerformIO)
import qualified Data.ByteString.Internal as S
import qualified Data.ByteString.Lazy.Internal as L
#endif

import qualified Blaze.ByteString.Builder.Internal as B
import Blaze.ByteString.Builder.Write (Write(..))
import qualified Blaze.ByteString.Builder.Word     as B
import Blaze.ByteString.Builder.Word (writeWord8)

import Criterion.Main

------------------------------------------------------------------------------
-- Benchmarks
------------------------------------------------------------------------------

main :: IO ()
main = defaultMain $ concat
    [ benchmark "mconcat . map fromWord8"
        myfromWord8s
        yourfromWord8s
        word8s
    ]
  where
    benchmark name putF builderF x =
        [ bench (name ++ " Put") $
            whnf (L.length . toLazyByteString2 . putF) x
        , bench (name ++ " Builder") $
            whnf (L.length . B.toLazyByteString . builderF) x
        ]

word8s :: [Word8]
word8s = take 100000 $ cycle [0..]
{-# NOINLINE word8s #-}

myfromWord8s :: [Word8] -> Put ()
myfromWord8s = putBuilder . mconcat . map fromWord8
{-# NOINLINE myfromWord8s #-}

yourfromWord8s :: [Word8] -> B.Builder
yourfromWord8s = mconcat . map B.fromWord8
{-# NOINLINE yourfromWord8s #-}


------------------------------------------------------------------------------
-- The Put type
------------------------------------------------------------------------------

data BufRange = BufRange {-# UNPACK #-} !(Ptr Word8) {-# UNPACK #-} !(Ptr Word8)

newtype Put a = Put {
    unPut :: forall r. (a -> PutStep r) -> PutStep r
  }

data PutSignal a =
    Done {-# UNPACK #-} !(Ptr Word8) a
  | BufferFull
      {-# UNPACK #-} !Int
      {-# UNPACK #-} !(Ptr Word8)
                     !(PutStep a)
  | InsertByteString
      {-# UNPACK #-} !(Ptr Word8) 
                     !S.ByteString
                     !(PutStep a)

type PutStep a =  UPtr -> UPtr -> IO (PutSignal a)

instance Monad Put where
  return x = Put $ \k -> k x
  {-# INLINE return #-}
  m >>= f  = Put $ \k -> unPut m (\x -> unPut (f x) k)
  {-# INLINE (>>=) #-}
  m >>  n  = Put $ \k -> unPut m (\_ -> unPut n k)
  {-# INLINE (>>) #-}

------------------------------------------------------------------------------
-- The Builder type with equal signals as the Put type
------------------------------------------------------------------------------

newtype Builder = Builder (forall r. PutStep r -> PutStep r)

instance Monoid Builder where
  mempty = Builder id
  {-# INLINE mempty #-}
  (Builder b1) `mappend` (Builder b2) = Builder $ b1 . b2
  {-# INLINE mappend #-}
  mconcat = foldr mappend mempty
  {-# INLINE mconcat #-}

putBuilder :: Builder -> Put ()
putBuilder (Builder build) = Put $ \k -> build (k ())

fromPut :: Put () -> Builder
fromPut (Put put) = Builder $ \k -> put (\_ -> k)

fromBuildStep :: (forall r. PutStep r -> BufRange -> IO (PutSignal r)) -> Builder
fromBuildStep step = Builder step'
  where
    step' k op ope = step k (BufRange (uptrToPtr op) (uptrToPtr ope))
{-# INLINE fromBuildStep #-}

callBuildStep :: PutStep a -> BufRange -> IO (PutSignal a)
callBuildStep k (BufRange op ope) = k (ptrToUPtr op) (ptrToUPtr ope)
{-# INLINE callBuildStep #-}

boxBuildStep :: PutStep a -> (BufRange -> IO (PutSignal a))
boxBuildStep step (BufRange op ope) = step (ptrToUPtr op) (ptrToUPtr ope)
{-# INLINE boxBuildStep #-}

unboxBuildStep :: (BufRange -> IO (PutSignal a)) -> PutStep a
unboxBuildStep step op ope = step (BufRange (uptrToPtr op) (uptrToPtr ope))
{-# INLINE unboxBuildStep #-}

fromWriteSingleton :: (a -> Write) -> a -> Builder
fromWriteSingleton write = 
    mkBuilder
  where
    mkBuilder x = fromBuildStep step
      where
        step k (BufRange pf pe)
          | pf `plusPtr` size <= pe = do
              io pf
              let !br' = BufRange (pf `plusPtr` size) pe
              callBuildStep k br'
          | otherwise               = 
              return $ BufferFull size pf (unboxBuildStep $ step k)
          where
            Write size io = write x
{-# INLINE fromWriteSingleton #-}

fromWord8 :: Word8 -> Builder
fromWord8 = fromWriteSingleton writeWord8
{-# INLINE fromWord8 #-}

------------------------------------------------------------------------------
-- More explicit implementation of running builders
------------------------------------------------------------------------------


data Buffer = Buffer {-# UNPACK #-} !(ForeignPtr Word8) -- underlying pinned array
                     {-# UNPACK #-} !(Ptr Word8)        -- beginning of slice
                     {-# UNPACK #-} !(Ptr Word8)        -- next free byte
                     {-# UNPACK #-} !(Ptr Word8)        -- first byte after buffer

allocBuffer :: Int -> IO Buffer
allocBuffer size = do
    fpbuf <- S.mallocByteString size
    let !pbuf = unsafeForeignPtrToPtr fpbuf
    return $! Buffer fpbuf pbuf pbuf (pbuf `plusPtr` size)

unsafeFreezeBuffer :: Buffer -> S.ByteString
unsafeFreezeBuffer (Buffer fpbuf p0 op _) = 
    S.PS fpbuf 0 (op `minusPtr` p0)

unsafeFreezeNonEmptyBuffer :: Buffer -> Maybe S.ByteString
unsafeFreezeNonEmptyBuffer (Buffer fpbuf p0 op _) 
  | p0 == op  = Nothing
  | otherwise = Just $ S.PS fpbuf 0 (op `minusPtr` p0)

nextSlice :: Int -> Buffer -> Maybe Buffer
nextSlice minSize (Buffer fpbuf _ op ope)
  | ope `minusPtr` op <= minSize = Nothing
  | otherwise                    = Just (Buffer fpbuf op op ope)

runPut :: Monad m 
       => (IO (PutSignal a) -> m (PutSignal a)) -- lifting of buildsteps
       -> (Int -> Buffer -> m Buffer) -- output function for a guaranteedly non-empty buffer, the returned buffer will be filled next
       -> (S.ByteString -> m ())    -- output function for guaranteedly non-empty bytestrings, that are inserted directly into the stream
       -> Put a                     -- put to execute
       -> Buffer                    -- initial buffer to be used
       -> m (a, Buffer)             -- result of put and remaining buffer
runPut liftIO outputBuf outputBS (Put put) =
    runStep (put $ (\x -> unboxBuildStep $ finalStep x))
  where
    finalStep x !(BufRange op _) = return $ Done op x

    runStep step buf@(Buffer fpbuf p0 op ope) = do
        let !br = BufRange op ope
        signal <- liftIO $ callBuildStep step br
        case signal of 
            Done op' x ->         -- put completed, buffer partially runSteped
                return (x, Buffer fpbuf p0 op' ope)

            BufferFull minSize op' nextStep -> do
                buf' <- outputBuf minSize (Buffer fpbuf p0 op' ope)
                runStep nextStep buf'

            InsertByteString op' bs nextStep
              | S.null bs ->   -- flushing of buffer required
                  outputBuf 1 (Buffer fpbuf p0 op' ope) >>= runStep nextStep
              | p0 == op' -> do -- no bytes written: just insert bytestring
                  outputBS bs
                  runStep nextStep buf
              | otherwise -> do   -- bytes written, insert buffer and bytestring
                  buf' <- outputBuf 1 (Buffer fpbuf p0 op' ope)
                  outputBS bs
                  runStep nextStep buf'
{-# INLINE runPut #-}
              
-- | A monad for lazily composing lazy bytestrings using continuations.
newtype LBSM a = LBSM { unLBSM :: (a, L.ByteString -> L.ByteString) }

instance Monad LBSM where
    return x                       = LBSM (x, id)
    (LBSM (x,k)) >>= f             = let LBSM (x',k') = f x in LBSM (x', k . k')
    (LBSM (_,k)) >> (LBSM (x',k')) = LBSM (x', k . k')

-- | Execute a put and return the written buffers as the chunks of a lazy
-- bytestring.
toLazyByteString2 :: Put a -> L.ByteString
toLazyByteString2 put = 
    k (bufToLBSCont (snd result) L.empty)
  where
    -- initial buffer
    buf0 = inlinePerformIO $ allocBuffer B.defaultBufferSize
    -- run put, but don't force result => we're lazy enough
    LBSM (result, k) = runPut liftIO outputBuf outputBS put buf0
    -- convert a buffer to a lazy bytestring continuation
    bufToLBSCont = maybe id L.Chunk . unsafeFreezeNonEmptyBuffer
    -- lifting an io putsignal to a lazy bytestring monad
    liftIO io = LBSM (inlinePerformIO io, id)
    -- add buffer as a chunk prepare allocation of new one
    outputBuf minSize buf = LBSM
        ( inlinePerformIO $ allocBuffer (max minSize B.defaultBufferSize)
        , bufToLBSCont buf )
    -- add bytestring directly as a chunk; exploits postcondition of runPut
    -- that bytestrings are non-empty
    outputBS bs = LBSM ((), L.Chunk bs)