Stack usage with a state monad

[Tomasz Zielonka]

On Tue, Dec 30, 2003 at 02:12:15PM +0000, Joe Thornber wrote:
> 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.

If you are using Hugs, try compiling your program with GHC (with -O2).

With GHC it seems to work, but it is still rather slow. After 4 minutes
of waiting a killed the process.

Correction: I had an environment option GHCRTS=-K64M, so it just took
more time before the stack exhausted.

I've optimised you program a bit and now it completes after 4 seconds
using only 2 megabytes of memory. After adding strictness annotations,
increasing sharing in the tree generated by buildTree the program still
was quite resource hungry, so I tried using an unboxed tuple (GHC's
extension) in the state monad - it helped a lot.

I am sorry, if I only confused you. My english is not great and time is
running. Got to go :)

Best regards,
Tom

{-# OPTIONS -fglasgow-exts #-}

module Main (module 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

newtype UnboxedState s a = UnboxedState { runUnboxedState :: s -> (# a, s #) }

instance Monad (UnboxedState s) where
	return a = UnboxedState $ \s -> (# a, s #)
	m >>= k  = UnboxedState $ \s ->
	    case runUnboxedState m s of
		(# a, s' #) -> runUnboxedState (k a) s'

instance MonadState s (UnboxedState s) where
	get   = UnboxedState $ \s -> (# s, s #)
	put s = UnboxedState $ \_ -> (# (), s #)

execUnboxedState m s = case runUnboxedState m s of
			(# _, s' #) -> s'

data Tree = Tree [Tree] | Leaf

buildTree :: Int -> Int -> Tree
buildTree order depth =
    head $ drop depth $ iterate (\t -> Tree (replicate order t)) Leaf

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

incCount :: UnboxedState Int ()
incCount = do {c <- get;
               put $! (c + 1);
              }

countLeaves2   :: Tree -> Int
countLeaves2 t = execUnboxedState (aux t) 0
    where
    aux (Tree xs) = mapM_ aux xs
    aux (Leaf) = incCount

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


-- 
.signature: Too many levels of symbolic links