Builder.hs

{-# LANGUAGE CPP #-}
-- |
-- Module      :  Pinch.Internal.Builder
-- Copyright   :  (c) Abhinav Gupta 2015
-- License     :  BSD3
--
-- Maintainer  :  Abhinav Gupta <mail@abhinavg.net>
-- Stability   :  experimental
--
-- This module provides a wrapper around @Data.ByteString.Builder@ that keeps
-- track of the final size of the generated ByteString.
module Pinch.Internal.Builder
    ( Builder
    , Build
    , run

    , int8
    , int16
    , int32
    , int64
    , double
    , byteString
    ) where

#if __GLASGOW_HASKELL__ < 709
import Control.Applicative
#endif

import Data.ByteString (ByteString)
import Data.Int
import Data.Monoid
import Prelude         hiding (length)

import qualified Data.ByteString         as B
import qualified Data.ByteString.Builder as BB

-- | Alias for a build that produces no results.
type Build = Builder ()

-- We're using big endian byte order for now. We can add little endian later
-- if needed.

-- | Used for building a long chain of ByteStrings.
data Builder a = B !a {-# UNPACK #-} !Int64 !BB.Builder
    -- TODO CPS?

instance Functor Builder where
    fmap f (B a l b) = B (f a) l b

instance Applicative Builder where
    pure a = B a 0 mempty
    B f l0 b0 <*> B a l1 b1 = B (f a) (l0 + l1) (b0 <> b1)

instance Monad Builder where
    (>>) = (*>)
    return = pure
    B a l0 b0 >>= k =
        let B b l1 b1 = k a
        in  B b (l0 + l1) (b0 <> b1)

-- | Returns the ByteString Builder for this build and its length.
run :: Build -> (Int64, BB.Builder)
run (B () l b) = (l, b)


-- | Writes a byte.
int8 :: Int8 -> Build
int8 = B () 1 . BB.int8


-- | Writes a 16-bit integer.
int16 :: Int16 -> Build
int16 = B () 2 . BB.int16BE


-- | Writes a 32-bit integer.
int32 :: Int32 -> Build
int32 = B () 4 . BB.int32BE


-- | Writes a 64-bit integer.
int64 :: Int64 -> Build
int64 = B () 8 . BB.int64BE


-- | Writes a 64-bit double.
double :: Double -> Build
double = B () 8 . BB.doubleBE


-- | Writes an arbitrary ByteString.
byteString :: ByteString -> Build
byteString bs = B () (fromIntegral $ B.length bs) (BB.byteString bs)