import Control.Monad
-- | State monad, where the state is an integer number.
data State a = S (Integer -> (a, Integer))
instance Monad State where
return x = S (\k -> (x, k))
S z >>= f = S (\k -> let (x, k') = z k
S z' = f x
in z' k')
get :: State Integer
get = S (\k -> (k, k))
set :: Integer -> State ()
set k' = S (\k -> ((), k'))
run :: Integer -> State a -> a
run k (S z) = fst $ z k
-- | Boilerplate code to make 'State' a legal monad.
instance Functor State where
fmap = liftM
instance Applicative State where
pure = return
(<*>) = ap
-- | A polymorphic n-ary tree data structure.
data Tree a = Node a [Tree a]
deriving Show
-- | Function that annotates each node of a tree with the order in which
-- | a DFS traversal would visit it.
dfs :: Tree a -> Tree (a, Integer)
dfs = run 1 . aux
where aux :: Tree a -> State (Tree (a, Integer))
aux (Node x ts) = do
k <- get
set $ k+1
ts' <- auxs ts
return $ Node (x, k) ts'
auxs :: [Tree a] -> State [Tree (a, Integer)]
auxs [] = return []
auxs (t : ts) = do
t' <- aux t
ts' <- auxs ts
return $ t' : ts'
-- | Examples
t1 = Node 1 [ Node 2 [ Node 3 []
, Node 4 []
]
, Node 5 [ Node 6 [] ]
]
t2 = Node 'a' [ Node 'b' []
, Node 'c' [ Node 'e' []
, Node 'f' []
]
, Node 'd' []
]
main = do
print (dfs t1)
print (dfs t2)