{-# LANGUAGE CPP, BangPatterns #-}

-- |
-- Module      : Blaze.ByteString.Builder.Internal.Poke
-- Copyright   : (c) 2010 Simon Meier
--               (c) 2010 Jasper van der Jeugt
-- License     : BSD3-style (see LICENSE)
--
-- Maintainer  : https://github.com/blaze-builder
-- Stability   : stable
-- Portability : tested on GHC only
--
-- A general and efficient write type that allows for the easy construction of
-- builders for (smallish) bounded size writes to a buffer.
--
-- FIXME: Improve documentation.
--
module Blaze.ByteString.Builder.Internal.Write (
  -- * Poking a buffer
    Poke(..)
  , pokeN

  -- * Writing to abuffer
  , Write(..)
  , runWrite
  , getBound
  , getBound'
  , getPoke

  , exactWrite
  , boundedWrite
  , writeLiftIO
  , writeIf
  , writeEq
  , writeOrdering
  , writeOrd

  -- * Constructing builders from writes
  , fromWrite
  , fromWriteSingleton
  , fromWriteList

  -- * Writing 'Storable's
  , writeStorable
  , fromStorable
  , fromStorables

  ) where

import Foreign

import qualified Data.Foldable as F
import Control.Monad

import Data.ByteString.Builder.Internal

#if !MIN_VERSION_base(4,8,0)
import Data.Monoid (Monoid(..))
#endif
import Data.Semigroup (Semigroup(..))

------------------------------------------------------------------------------
-- Poking a buffer and writing to a buffer
------------------------------------------------------------------------------

-- Sadly GHC is not smart enough: code where we branch and each branch should
-- execute a few IO actions and then return a value cannot be taught to GHC. At
-- least not such that it returns the value of the branches unpacked.
--
-- Hmm.. at least he behaves much better for the Monoid instance of Write
-- than the one for Poke. Serializing UTF-8 chars gets a slowdown of a
-- factor 2 when 2 chars are composed. Perhaps I should try out the writeList
-- instances also, as they may be more sensitive to to much work per Char.
--

-- | Changing a sequence of bytes starting from the given pointer. 'Poke's are
-- the most primitive buffer manipulation. In most cases, you don't use the
-- explicitly but as part of a 'Write', which also tells how many bytes will
-- be changed at most.
newtype Poke =
    Poke { Poke -> Ptr Word8 -> IO (Ptr Word8)
runPoke :: Ptr Word8 -> IO (Ptr Word8) }

-- | A write of a bounded number of bytes.
--
-- When defining a function @write :: a -> Write@ for some @a@, then it is
-- important to ensure that the bound on the number of bytes written is
-- data-independent. Formally,
--
--  @ forall x y. getBound (write x) = getBound (write y) @
--
-- The idea is that this data-independent bound is specified such that the
-- compiler can optimize the check, if there are enough free bytes in the buffer,
-- to a single subtraction between the pointer to the next free byte and the
-- pointer to the end of the buffer with this constant bound of the maximal
-- number of bytes to be written.
--
data Write = Write {-# UNPACK #-} !Int Poke

-- | Extract the 'Poke' action of a write.
{-# INLINE getPoke #-}
getPoke :: Write -> Poke
getPoke :: Write -> Poke
getPoke (Write Int
_ Poke
wio) = Poke
wio

-- | Run the 'Poke' action of a write.
{-# INLINE runWrite #-}
runWrite :: Write -> Ptr Word8 -> IO (Ptr Word8)
runWrite :: Write -> Ptr Word8 -> IO (Ptr Word8)
runWrite = Poke -> Ptr Word8 -> IO (Ptr Word8)
runPoke forall b c a. (b -> c) -> (a -> b) -> a -> c
. Write -> Poke
getPoke

-- | Extract the maximal number of bytes that this write could write.
{-# INLINE getBound #-}
getBound :: Write -> Int
getBound :: Write -> Int
getBound (Write Int
bound Poke
_) = Int
bound

-- | Extract the maximal number of bytes that this write could write in any
-- case. Assumes that the bound of the write is data-independent.
{-# INLINE getBound' #-}
getBound' :: String             -- ^ Name of caller: for debugging purposes.
          -> (a -> Write)
          -> Int
getBound' :: forall a. String -> (a -> Write) -> Int
getBound' String
msg a -> Write
write =
    Write -> Int
getBound forall a b. (a -> b) -> a -> b
$ a -> Write
write forall a b. (a -> b) -> a -> b
$ forall a. HasCallStack => String -> a
error forall a b. (a -> b) -> a -> b
$
    String
"getBound' called from " forall a. [a] -> [a] -> [a]
++ String
msg forall a. [a] -> [a] -> [a]
++ String
": write bound is not data-independent."

instance Semigroup Poke where
  {-# INLINE (<>) #-}
  (Poke Ptr Word8 -> IO (Ptr Word8)
po1) <> :: Poke -> Poke -> Poke
<> (Poke Ptr Word8 -> IO (Ptr Word8)
po2) = (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke forall a b. (a -> b) -> a -> b
$ Ptr Word8 -> IO (Ptr Word8)
po1 forall (m :: * -> *) a b c.
Monad m =>
(a -> m b) -> (b -> m c) -> a -> m c
>=> Ptr Word8 -> IO (Ptr Word8)
po2

  {-# INLINE sconcat #-}
  sconcat :: NonEmpty Poke -> Poke
sconcat = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr forall a. Semigroup a => a -> a -> a
(<>) forall a. Monoid a => a
mempty

instance Monoid Poke where
  {-# INLINE mempty #-}
  mempty :: Poke
mempty = (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke forall a b. (a -> b) -> a -> b
$ forall (m :: * -> *) a. Monad m => a -> m a
return

#if !(MIN_VERSION_base(4,11,0))
  {-# INLINE mappend #-}
  mappend = (<>)

  {-# INLINE mconcat #-}
  mconcat = F.foldr mappend mempty
#endif

instance Semigroup Write where
  {-# INLINE (<>) #-}
  (Write Int
bound1 Poke
w1) <> :: Write -> Write -> Write
<> (Write Int
bound2 Poke
w2) =
    Int -> Poke -> Write
Write (Int
bound1 forall a. Num a => a -> a -> a
+ Int
bound2) (Poke
w1 forall a. Semigroup a => a -> a -> a
<> Poke
w2)

  {-# INLINE sconcat #-}
  sconcat :: NonEmpty Write -> Write
sconcat = forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
F.foldr forall a. Semigroup a => a -> a -> a
(<>) forall a. Monoid a => a
mempty

instance Monoid Write where
  {-# INLINE mempty #-}
  mempty :: Write
mempty = Int -> Poke -> Write
Write Int
0 forall a. Monoid a => a
mempty

#if !(MIN_VERSION_base(4,11,0))
  {-# INLINE mappend #-}
  mappend = (<>)

  {-# INLINE mconcat #-}
  mconcat = F.foldr mappend mempty
#endif

-- | @pokeN size io@ creates a write that denotes the writing of @size@ bytes
-- to a buffer using the IO action @io@. Note that @io@ MUST write EXACTLY @size@
-- bytes to the buffer!
{-# INLINE pokeN #-}
pokeN :: Int
       -> (Ptr Word8 -> IO ()) -> Poke
pokeN :: Int -> (Ptr Word8 -> IO ()) -> Poke
pokeN Int
size Ptr Word8 -> IO ()
io = (Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke forall a b. (a -> b) -> a -> b
$ \Ptr Word8
op -> Ptr Word8 -> IO ()
io Ptr Word8
op forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> (forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$! (Ptr Word8
op forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
size))


-- | @exactWrite size io@ creates a bounded write that can later be converted to
-- a builder that writes exactly @size@ bytes. Note that @io@ MUST write
-- EXACTLY @size@ bytes to the buffer!
{-# INLINE exactWrite #-}
exactWrite :: Int
           -> (Ptr Word8 -> IO ())
           -> Write
exactWrite :: Int -> (Ptr Word8 -> IO ()) -> Write
exactWrite Int
size Ptr Word8 -> IO ()
io = Int -> Poke -> Write
Write Int
size (Int -> (Ptr Word8 -> IO ()) -> Poke
pokeN Int
size Ptr Word8 -> IO ()
io)

-- | @boundedWrite size write@ creates a bounded write from a @write@ that does
-- not write more than @size@ bytes.
{-# INLINE boundedWrite #-}
boundedWrite :: Int -> Poke -> Write
boundedWrite :: Int -> Poke -> Write
boundedWrite = Int -> Poke -> Write
Write

-- | @writeLiftIO io write@ creates a write executes the @io@ action to compute
-- the value that is then written.
{-# INLINE writeLiftIO #-}
writeLiftIO :: (a -> Write) -> IO a -> Write
writeLiftIO :: forall a. (a -> Write) -> IO a -> Write
writeLiftIO a -> Write
write IO a
io =
    Int -> Poke -> Write
Write (forall a. String -> (a -> Write) -> Int
getBound' String
"writeLiftIO" a -> Write
write)
          ((Ptr Word8 -> IO (Ptr Word8)) -> Poke
Poke forall a b. (a -> b) -> a -> b
$ \Ptr Word8
pf -> do a
x <- IO a
io; Write -> Ptr Word8 -> IO (Ptr Word8)
runWrite (a -> Write
write a
x) Ptr Word8
pf)

-- | @writeIf p wTrue wFalse x@ creates a 'Write' with a 'Poke' equal to @wTrue
-- x@, if @p x@ and equal to @wFalse x@ otherwise. The bound of this new
-- 'Write' is the maximum of the bounds for either 'Write'. This yields a data
-- independent bound, if the bound for @wTrue@ and @wFalse@ is already data
-- independent.
{-# INLINE writeIf #-}
writeIf :: (a -> Bool) -> (a -> Write) -> (a -> Write) -> (a -> Write)
writeIf :: forall a. (a -> Bool) -> (a -> Write) -> (a -> Write) -> a -> Write
writeIf a -> Bool
p a -> Write
wTrue a -> Write
wFalse a
x =
    Int -> Poke -> Write
boundedWrite (forall a. Ord a => a -> a -> a
max (Write -> Int
getBound forall a b. (a -> b) -> a -> b
$ a -> Write
wTrue a
x) (Write -> Int
getBound forall a b. (a -> b) -> a -> b
$ a -> Write
wFalse a
x))
                 (if a -> Bool
p a
x then Write -> Poke
getPoke forall a b. (a -> b) -> a -> b
$ a -> Write
wTrue a
x else Write -> Poke
getPoke forall a b. (a -> b) -> a -> b
$ a -> Write
wFalse a
x)

-- | Compare the value to a test value and use the first write action for the
-- equal case and the second write action for the non-equal case.
{-# INLINE writeEq #-}
writeEq :: Eq a => a -> (a -> Write) -> (a -> Write) -> (a -> Write)
writeEq :: forall a. Eq a => a -> (a -> Write) -> (a -> Write) -> a -> Write
writeEq a
test = forall a. (a -> Bool) -> (a -> Write) -> (a -> Write) -> a -> Write
writeIf (a
test forall a. Eq a => a -> a -> Bool
==)

-- | TODO: Test this. It might well be too difficult to use.
--   FIXME: Better name required!
{-# INLINE writeOrdering #-}
writeOrdering :: (a -> Ordering)
              -> (a -> Write) -> (a -> Write) -> (a -> Write)
              -> (a -> Write)
writeOrdering :: forall a.
(a -> Ordering)
-> (a -> Write) -> (a -> Write) -> (a -> Write) -> a -> Write
writeOrdering a -> Ordering
ord a -> Write
wLT a -> Write
wEQ a -> Write
wGT a
x =
    Int -> Poke -> Write
boundedWrite Int
bound (case a -> Ordering
ord a
x of Ordering
LT -> Write -> Poke
getPoke forall a b. (a -> b) -> a -> b
$ a -> Write
wLT a
x;
                                      Ordering
EQ -> Write -> Poke
getPoke forall a b. (a -> b) -> a -> b
$ a -> Write
wEQ a
x;
                                      Ordering
GT -> Write -> Poke
getPoke forall a b. (a -> b) -> a -> b
$ a -> Write
wGT a
x)
  where
    bound :: Int
bound = forall a. Ord a => a -> a -> a
max (Write -> Int
getBound forall a b. (a -> b) -> a -> b
$ a -> Write
wLT a
x) (forall a. Ord a => a -> a -> a
max (Write -> Int
getBound forall a b. (a -> b) -> a -> b
$ a -> Write
wEQ a
x) (Write -> Int
getBound forall a b. (a -> b) -> a -> b
$ a -> Write
wGT a
x))

-- | A write combinator useful to build decision trees for deciding what value
-- to write with a constant bound on the maximal number of bytes written.
{-# INLINE writeOrd #-}
writeOrd :: Ord a
       => a
       -> (a -> Write) -> (a -> Write) -> (a -> Write)
       -> (a -> Write)
writeOrd :: forall a.
Ord a =>
a -> (a -> Write) -> (a -> Write) -> (a -> Write) -> a -> Write
writeOrd a
test = forall a.
(a -> Ordering)
-> (a -> Write) -> (a -> Write) -> (a -> Write) -> a -> Write
writeOrdering (forall a. Ord a => a -> a -> Ordering
`compare` a
test)

-- | Create a builder that execute a single 'Write'.
{-# INLINE fromWrite #-}
fromWrite :: Write -> Builder
fromWrite :: Write -> Builder
fromWrite (Write Int
maxSize Poke
wio) =
    (forall r. BuildStep r -> BuildStep r) -> Builder
builder forall r. BuildStep r -> BuildStep r
step
  where
    step :: (BufferRange -> IO (BuildSignal a))
-> BufferRange -> IO (BuildSignal a)
step BufferRange -> IO (BuildSignal a)
k (BufferRange Ptr Word8
op Ptr Word8
ope)
      | Ptr Word8
op forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
maxSize forall a. Ord a => a -> a -> Bool
<= Ptr Word8
ope = do
          Ptr Word8
op' <- Poke -> Ptr Word8 -> IO (Ptr Word8)
runPoke Poke
wio Ptr Word8
op
          let !br' :: BufferRange
br' = Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
op' Ptr Word8
ope
          BufferRange -> IO (BuildSignal a)
k BufferRange
br'
      | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
bufferFull Int
maxSize Ptr Word8
op ((BufferRange -> IO (BuildSignal a))
-> BufferRange -> IO (BuildSignal a)
step BufferRange -> IO (BuildSignal a)
k)

{-# INLINE fromWriteSingleton #-}
fromWriteSingleton :: (a -> Write) -> (a -> Builder)
fromWriteSingleton :: forall a. (a -> Write) -> a -> Builder
fromWriteSingleton a -> Write
write =
    a -> Builder
mkBuilder
  where
    mkBuilder :: a -> Builder
mkBuilder a
x = (forall r. BuildStep r -> BuildStep r) -> Builder
builder forall r. BuildStep r -> BuildStep r
step
      where
        step :: (BufferRange -> IO (BuildSignal a))
-> BufferRange -> IO (BuildSignal a)
step BufferRange -> IO (BuildSignal a)
k (BufferRange Ptr Word8
op Ptr Word8
ope)
          | Ptr Word8
op forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
maxSize forall a. Ord a => a -> a -> Bool
<= Ptr Word8
ope = do
              Ptr Word8
op' <- Poke -> Ptr Word8 -> IO (Ptr Word8)
runPoke Poke
wio Ptr Word8
op
              let !br' :: BufferRange
br' = Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
op' Ptr Word8
ope
              BufferRange -> IO (BuildSignal a)
k BufferRange
br'
          | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
bufferFull Int
maxSize Ptr Word8
op ((BufferRange -> IO (BuildSignal a))
-> BufferRange -> IO (BuildSignal a)
step BufferRange -> IO (BuildSignal a)
k)
          where
            Write Int
maxSize Poke
wio = a -> Write
write a
x


-- | Construct a 'Builder' writing a list of data one element at a time.
fromWriteList :: (a -> Write) -> [a] -> Builder
fromWriteList :: forall a. (a -> Write) -> [a] -> Builder
fromWriteList a -> Write
write =
    [a] -> Builder
makeBuilder
  where
    makeBuilder :: [a] -> Builder
makeBuilder [a]
xs0 = (forall r. BuildStep r -> BuildStep r) -> Builder
builder forall a b. (a -> b) -> a -> b
$ forall {a}.
[a]
-> (BufferRange -> IO (BuildSignal a))
-> BufferRange
-> IO (BuildSignal a)
step [a]
xs0
      where
        step :: [a]
-> (BufferRange -> IO (BuildSignal a))
-> BufferRange
-> IO (BuildSignal a)
step [a]
xs1 BufferRange -> IO (BuildSignal a)
k !(BufferRange Ptr Word8
op0 Ptr Word8
ope0) = [a] -> Ptr Word8 -> IO (BuildSignal a)
go [a]
xs1 Ptr Word8
op0
          where
            go :: [a] -> Ptr Word8 -> IO (BuildSignal a)
go [] !Ptr Word8
op = do
               let !br' :: BufferRange
br' = Ptr Word8 -> Ptr Word8 -> BufferRange
BufferRange Ptr Word8
op Ptr Word8
ope0
               BufferRange -> IO (BuildSignal a)
k BufferRange
br'

            go xs :: [a]
xs@(a
x':[a]
xs') !Ptr Word8
op
              | Ptr Word8
op forall a b. Ptr a -> Int -> Ptr b
`plusPtr` Int
maxSize forall a. Ord a => a -> a -> Bool
<= Ptr Word8
ope0 = do
                  !Ptr Word8
op' <- Poke -> Ptr Word8 -> IO (Ptr Word8)
runPoke Poke
wio Ptr Word8
op
                  [a] -> Ptr Word8 -> IO (BuildSignal a)
go [a]
xs' Ptr Word8
op'
              | Bool
otherwise = forall (m :: * -> *) a. Monad m => a -> m a
return forall a b. (a -> b) -> a -> b
$ forall a. Int -> Ptr Word8 -> BuildStep a -> BuildSignal a
bufferFull Int
maxSize Ptr Word8
op ([a]
-> (BufferRange -> IO (BuildSignal a))
-> BufferRange
-> IO (BuildSignal a)
step [a]
xs BufferRange -> IO (BuildSignal a)
k)
              where
                Write Int
maxSize Poke
wio = a -> Write
write a
x'
{-# INLINE fromWriteList #-}



------------------------------------------------------------------------------
-- Writing storables
------------------------------------------------------------------------------


-- | Write a storable value.
{-# INLINE writeStorable #-}
writeStorable :: Storable a => a -> Write
writeStorable :: forall a. Storable a => a -> Write
writeStorable a
x = Int -> (Ptr Word8 -> IO ()) -> Write
exactWrite (forall a. Storable a => a -> Int
sizeOf a
x) (\Ptr Word8
op -> forall a. Storable a => Ptr a -> a -> IO ()
poke (forall a b. Ptr a -> Ptr b
castPtr Ptr Word8
op) a
x)

-- | A builder that serializes a storable value. No alignment is done.
{-# INLINE fromStorable #-}
fromStorable :: Storable a => a -> Builder
fromStorable :: forall a. Storable a => a -> Builder
fromStorable = forall a. (a -> Write) -> a -> Builder
fromWriteSingleton forall a. Storable a => a -> Write
writeStorable

-- | A builder that serializes a list of storable values by writing them
-- consecutively. No alignment is done. Parsing information needs to be
-- provided externally.
fromStorables :: Storable a => [a] -> Builder
fromStorables :: forall a. Storable a => [a] -> Builder
fromStorables = forall a. (a -> Write) -> [a] -> Builder
fromWriteList forall a. Storable a => a -> Write
writeStorable