{-# LANGUAGE CPP #-}

import qualified Data.ByteString.Lazy as BL
import qualified Data.Binary.Put as P
 
import Control.Monad.State
import Control.Monad.Writer

-- I'm trying to wrap the PutM monad from Bada.Binary.Put
-- module with another monad so that later I can ask it questions
-- like "what is the current offset" and "how many bytes is this
-- going to write", where my intention is to use StateT
-- and SumT (defined below) respectively for that.
-- 
-- Problem is that both of these transformers cause a huge
-- memory consumption. Now I kind of have two questions:
--
-- 1.) I understand this will probably have to do something
--     with lazyness, but what exactly is going on inside?
--
-- 2.) Can it be fixed?

-- To test it, redefine the USE_TRANSFORMER macro to one
-- of these values:
--
--   1: IdentityT    Works without a problem, but is useless.
--   2: SumT         Consumes a lot of memory
--   3: StateT       Consumes a lot of memory

#define USE_TRANSFORMER 3

-- -------------------------------------------------------------------
#if USE_TRANSFORMER == 1

newtype IdentityT m a = IdentityT { runIdT :: m a }

instance (Monad m) => Monad (IdentityT m) where
  return a = IdentityT $ return a
  ma >>= f = IdentityT $ runIdT ma >>= runIdT . f
  ma >> mb = IdentityT $ runIdT ma >> runIdT mb
  
instance MonadTrans IdentityT where 
  lift ma = IdentityT ma

type Out = IdentityT P.PutM ()

writeToFile :: String -> Out -> IO ()
writeToFile path out = BL.writeFile path (P.runPut $ runIdT out >> return ())

-- -------------------------------------------------------------------
#elif USE_TRANSFORMER == 2

instance Monoid Integer where
  mappend = (+)
  mempty = 0

type Out = WriterT Integer P.PutM ()

writeToFile :: String -> Out -> IO ()
writeToFile path out = do
  BL.writeFile path $ P.runPut $ runWriterT out >> return ()

-- -------------------------------------------------------------------
#elif USE_TRANSFORMER == 3

type Out = StateT Integer P.PutM ()

writeToFile :: String -> Out -> IO ()
writeToFile path out = BL.writeFile path $ P.runPut $ runStateT out 0 >> return ()

#endif
-- -------------------------------------------------------------------
-- Test data.

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

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 = lift $ P.putWord32le n
putInt8  n = lift $ P.putWord8 n

putTree :: Tree -> Out
putTree (Node childs) = do
  putInt8 123
  mapM_ putTree childs

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

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

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


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