How do you save a tree data structure to binary files in non linear fashion using Haskell

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
{-# LANGUAGE CPP,Rank2Types #-}

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

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

type Offset = Int
type Size   = Int

-- -------------------------------------------------------------------
-- A State transformer using continuations as suggested by nominolo
-- here: http://stackoverflow.com/questions/5019655/why-wrapping-the-data-binary-put-monad-creates-a-memory-leak-part-2
--

-- The S transformer is used to trac current offset.
-- The WriterT is used to trac child sizes.

newtype S s m a = S { unS :: forall r . (s -> a -> m r) -> s -> m r }

instance Monad m => Monad (S s m) where
  return a = S $ \f s -> f s a
  ma >>= f = S $ \fb s -> unS ma (\s' a -> unS (f a) fb s') s

instance MonadTrans (S s) where
  lift ma = S $ \f s -> ma >>= f s

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

type MyPut = S Offset (WriterT Size P.PutM)

peal_1 put = unS put (\o -> return) 0
peal_2 wri = runWriterT wri

writeToFile :: String -> MyPut () -> IO ()
writeToFile path put =
  BL.writeFile path $ P.runPut $ peal_2 (peal_1 put) >> return ()

getSize :: MyPut a -> MyPut Int
getSize o = lift $ listen (peal_1 o) >>= return . snd

getCurrentOffset :: MyPut Int
getCurrentOffset = S $ \f o -> f o o

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

putTree :: Tree -> MyPut ()
putTree (Node childs) = do
  curOffset <- getCurrentOffset
  sizes <- mapM (getSize . putTree) childs
  let offsetTableSize = 4 * length sizes -- Offset is writen in 4 bytes
      offsets = scanl (+) (curOffset+offsetTableSize) sizes
  mapM_ putInt32 $ take (length sizes) offsets
  mapM_ putTree childs

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

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

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


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