Stack usage with a state monad

[Joe Thornber]

Hi,

I was wondering if anyone could give me some help with this problem ?

I'm trying to hold some state in a StateMonad whilst I iterate over a
large tree, and finding that I'm running out of stack space very
quickly.  The simplified program below exhibits the same problem.

This is the first time I've hit space problems in Haskell, I hope
judicial use of 'seq' or '$!' would be enough to fix it, but I don't
know where to start.

Any ideas as to what I'm doing wrong would be much appreciated.

Thanks,

- Joe




module Main (main) where

-- Program to count the leaf nodes in a rose tree.  Written to try and
-- reproduce a stack space leak present in a larger program.

-- How can I use a state monad to count the leaves without eating all
-- the stack ?

import Control.Monad.State

data Tree = Tree [Tree] | Leaf

buildTree :: Int -> Int -> Tree
buildTree order = buildTree'
    where
    buildTree' 0 = Leaf
    buildTree' depth = Tree $ map (buildTree') $ take order $ repeat (depth - 1)

countLeaves1 :: Tree -> Int
countLeaves1 (Tree xs) = sum $ map (countLeaves1) xs
countLeaves1 (Leaf) = 1

incCount :: State Int ()
incCount = do {c <- get;
	       put (c + 1);
	       return ();
	      }

countLeaves2   :: Tree -> Int
countLeaves2 t = execState (aux t) 0
    where
    aux :: Tree -> State Int ()
    aux (Tree xs) = foldr1 (>>) $ map (aux) xs
    aux (Leaf) = incCount

main :: IO ()B
main = print $ countLeaves2 $ buildTree 15 6