Why changing the Data.Binary.Put monad into a transformer creates a memory leak?

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
{-# 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"


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