StateExample.hs

import Data.Tree
import Data.Traversable
import Control.Monad.State hiding (mapM)

import Prelude hiding (mapM)

test :: (Tree Integer, Integer)
test = markOrdering (Node undefined [Node undefined [Node undefined []], Node undefined []]) 0

markOrdering :: (Num t1) => Tree t -> t1 -> (Tree t1, t1)
markOrdering (Node _ ts) count = (Node count ts', count')
  where 
    (ts', count') = chain markOrdering (count + 1) ts

chain :: (t -> t1 -> (a, t1)) -> t1 -> [t] -> ([a], t1)
chain f count [] = ([], count)
chain f count (x:xs) = (a:others, result) 
  where
    (a, count') = f x count
    (others, result) = chain f count' xs

test2 :: (Tree Integer, Integer)
test2 = runState (markOrderingM (Node undefined [Node undefined [Node undefined []], Node undefined []])) 0

markOrderingM :: (Num a) => Tree t -> State a (Tree a)
markOrderingM (Node _ ts) = State markOrderingM' 
  where 
    markOrderingM' count = (Node count ts', count')
      where
        (ts', count') = runState (mapM markOrderingM ts) (count + 1)


test3 :: (Tree Integer, Integer)
test3 = runState (markOrderingM2 (Node undefined [Node undefined [Node undefined []], Node undefined []])) 0

markOrderingM2 :: (MonadState s m, Num s) => Tree t -> m (Tree s)
markOrderingM2 (Node _ ts) = do
  count <- get
  put (count + 1)
  ts' <- mapM markOrderingM2 ts
  return $ Node count ts'

markOrderingM3 :: (MonadState a m, Num a, Traversable t) => t b -> m (t a)
markOrderingM3 t = mapM (const $ get >>= \x -> put (x+1) >> return x) t

-- really should be in standard libs:
getAndModify :: (MonadState s m) => (s -> s) -> m s
getAndModify f = get >>= \x -> put (f x) >> return x

test4 :: (Tree Integer, Integer)
test4 = runState (markOrderingM4 (Node undefined [Node undefined [Node undefined []], Node undefined []])) 0

markOrderingM4 :: (Num s, MonadState s m, Traversable t) => t b -> m (t s)
markOrderingM4 t = mapM (const $ getAndModify (+1)) t