{-# LANGUAGE CPP,Rank2Types,BangPatterns #-}

-- ghc -auto-all -fforce-recomp -O2 --make test5.hs \
-- && ./test5 +RTS -hT \
-- && hp2ps -c test5.hp \
-- && okular test5.ps

import Control.Monad.Writer

import Data.Monoid
import Data.Binary.Builder (Builder, toLazyByteString)
import qualified Data.Binary.Builder as B

import Data.Word
import qualified Data.ByteString.Lazy as L
import Control.Monad.Identity

-- Redefine USE_CASE to test different approaches
--   0 for no changes, just using Data.Binary.Put.
--   1 for adding Identity monad into D.B.Put.
--   2 is just a test if continuations would help but had no luck.
#define USE_CASE 1

-- -------------------------------------------------------------------
#if USE_CASE == 0
-- The original code as it is in Data.Binary.Put, works without
-- creating a space leak.

data PairS a = PairS a {-# UNPACK #-}!Builder

sndS :: PairS a -> Builder
sndS (PairS _ b) = b

newtype PutM a = Put { unPut :: PairS a }

type Put = PutM ()

instance Monad PutM where
    return a = Put $ PairS a mempty
    {-# INLINE return #-}

    m >>= k  = Put $
        let PairS a w  = unPut m
            PairS b w' = unPut (k a)
        in PairS b (w `mappend` w')
    {-# INLINE (>>=) #-}

    m >> k  = Put $
        let PairS _ w  = unPut m
            PairS b w' = unPut k
        in PairS b (w `mappend` w')
    {-# INLINE (>>) #-}

tell' :: Builder -> Put
tell' b = Put $ PairS () b
{-# INLINE tell' #-}

runPut :: Put -> L.ByteString
runPut = toLazyByteString . sndS . unPut
{-# INLINE runPut #-}

-- -------------------------------------------------------------------
#elif USE_CASE == 1
-- Here is almost exactly the same code as above with
-- the only difference that the PairS inside the Put
-- structure is wrapped inside Identity monad.
-- It hogs memory quite a lot.

data PairS a = PairS a {-# UNPACK #-}!Builder

sndS :: PairS a -> Builder
sndS (PairS _ !b) = b

newtype PutM a = Put { unPut :: Identity (PairS a) }

type Put = PutM ()

instance Monad PutM where
    return a = Put $! return $! PairS a mempty
    {-# INLINE return #-}

    m >>= k  = Put $!
        do PairS a w  <- unPut m
           PairS b w' <- unPut (k a)
           return $! PairS b $! (w `mappend` w')
    {-# INLINE (>>=) #-}

    m >> k  = Put $!
        do PairS _ w  <- unPut m
           PairS b w' <- unPut k
           return $! PairS b $! (w `mappend` w')
    {-# INLINE (>>) #-}

tell' :: Builder -> Put
tell' b = Put $! return $! PairS () b

runPut :: Put -> L.ByteString
runPut = toLazyByteString . sndS . runIdentity . unPut

-- -------------------------------------------------------------------
#elif USE_CASE == 2
-- It is said that sometimes using continuation monad
-- helps in forcing evaluation evaluation, unfortunately 
-- this is not the case, at least not the way I tried
-- it:

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

type Put = PutM ()

instance Monad PutM where
  return a = Put $! \f b -> f b a
  ma >>= f = Put $! \fb s -> unPut ma (\s' a -> unPut (f a) fb s') s

tell' :: Builder -> Put
tell' b = Put $! \f b' -> f (b `mappend` b') ()

runPut :: Put -> L.ByteString
runPut p = toLazyByteString $ unPut p (\b _ -> b) mempty

#endif

------------------------------------------------------------------------

putWord8 :: Word8 -> Put
putWord8 = tell' . B.singleton
{-# INLINE putWord8 #-}

putWord32le :: Word32 -> Put
putWord32le = tell' . B.putWord32le
{-# INLINE putWord32le #-}

-- -------------------------------------------------------------------
-- Test data.

data Tree = Node [Tree] | Leaf [Int] deriving Show

-- Approximate size in memory (ignoring laziness) I think is:
-- 101 * 4^9 * sizeof(Int) + 1/3 * 4^9 * sizeof(Node)

makeTree :: Tree
makeTree = makeTree' 9
  where makeTree' 0 = Leaf [0..100]
        makeTree' n = Node [ makeTree' $ n - 1
                           , makeTree' $ n - 1
                           , makeTree' $ n - 1
                           , makeTree' $ n - 1 ]

-- -------------------------------------------------------------------

putInt32 n = putWord32le $ fromIntegral n
putInt8  n = putWord8 $ fromIntegral n

putTree :: Tree -> PutM ()
putTree (Node childs) = do
  mapM_ putInt32 [1,2,3,4]
  mapM_ putTree childs

putTree (Leaf nums) = do
  mapM_ (putInt32) nums

-- -------------------------------------------------------------------

writeToFile :: String -> PutM () -> IO ()
writeToFile path put =
  L.writeFile path $ runPut $ put >> return ()

main = do
  putStrLn "begin"
  writeToFile "test-output.bin" $ putTree makeTree
  putStrLn "end"


-- -------------------------------------------------------------------