Re: Stack usage with a state monad

30 Dec 2003

      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

Re: Stack usage with a state monad

Tomasz Zielonka