import Control.Applicative
import Control.Monad (ap, liftM)
data State s a = State (s -> (a, s))
instance Monad (State s) where
return x = State (\s -> (x, s))
State f >>= g = State (\s -> let (a, s') = f s
State h = g a
in h s')
instance Functor (State s) where
fmap = liftM
instance Applicative (State s) where
pure = return
(<*>) = ap
get :: State s s
get = State (\s -> (s, s))
modify :: (s -> s) -> State s ()
modify f = State (\s -> ((), f s))
run :: s -> State s a -> a
run s (State f) = fst $ f s
data Tree a = Node a [Tree a]
deriving Show
t = Node 'a' [ Node 'b' []
, Node 'c' [ Node 'f' []
, Node 'g' []
]
, Node 'd' []
, Node 'e' [ Node 'h' []
]
]
dfsM :: Tree a -> Tree (a, Int)
dfsM = run 1 . aux
where aux :: Tree a -> State Int (Tree (a, Int))
aux (Node x ts) = do id <- get
modify (\s -> id+1)
ts' <- auxs ts
return $ Node (x, id) ts'
auxs :: [Tree a] -> State Int [Tree (a, Int)]
auxs [] = return []
auxs (t : ts) = do t' <- aux t
ts' <- auxs ts
return $ t' : ts'
main = do
putStrLn "Warming up with the state monad..."
let s :: Int
s = 1
let x = run s $ do one <- get
modify (+ 1)
two <- get
modify (* 3)
six <- get
modify (\n -> n * (n+1))
result <- get
return (one, two, six, result)
print x
putStrLn "Original tree:"
print t
putStrLn "Tree annotated with DFS numbering:"
print $ dfsM t