How to decorate a Tree in Haskell

Question

I want to tag each element of a tree with a different value (Int, for example sake). I managed to do this but the code is ugly as a beast and I don\'t know how to work with

Answer 1

I've modified your types slightly. Study this code carefully:

import Control.Monad.State

-- It's better not to use a pair as the argument of the constructor    
data Tree a = Tree a [Tree a] deriving Show

-- We typically want to put the Tree argument last; it makes it
-- easier to compose tree functions.  
--
-- Also, the Enum class is what you want here instead of numbers; 
-- you want a "give me the next tag" operation, which is the succ
-- method from Enum.  (For Int, succ is (+1).)
tag :: Enum t => t -> Tree a -> Tree (a, t)
tag init tree = 
    -- tagStep is where the action happens.  This just gets the ball
    -- rolling.
    evalState (tagStep tree) init

-- This is one monadic "step" of the calculation.  It assumes that
-- it has access to the current tag value implicitcly.  I'll 
-- annotate it in the comments.
tagStep :: Enum t => Tree a -> State t (Tree (a, t))
tagStep (Tree a subtrees) = 
    do -- First, recurse into the subtrees.  mapM is a utility function
       -- for executing a monadic action (like tagStep) on a list of
       -- elements, and producing the list of results.
       subtrees' <- mapM tagStep subtrees  

       -- The monadic action "get" accesses the implicit state parameter
       -- in the State monad.  The variable tag gets the value.
       tag <- get 

       -- The monadic action `put` sets the implicit state parameter in
       -- the State monad.  The next get will see the value of succ tag
       -- (assuming no other puts in between).
       --
       -- Note that when we did mapM tagStep subtrees above, this will 
       -- have executed a get and a put (succ tag) for each subtree.           
       put (succ tag)

       return $ Tree (a, tag) subtrees'

---

EDIT: Same solution as above, but put through one round of refactoring into reusable pieces:

-- This function is not part of the solution, but it will help you 
-- understand mapTreeM below.
mapTree :: (a -> b) -> Tree a -> Tree b
mapTree fn (Tree a subtrees) = 
    let subtrees' = map (mapTree fn) subtrees
        a' = fn a
     in Tree a' subtrees'

-- Normally you'd write that function like this:
mapTree' fn (Tree a subtrees) = Tree (fn a) $ map (mapTree' fn) subtrees

-- But I wrote it out the long way to bring out the similarity to the 
-- following, which extracts the structure of the tagStep definition from 
-- the first solution above.    
mapTreeM :: Monad m => (a -> m b) -> Tree a -> m (Tree b)
mapTreeM action (Tree a subtrees) =
    do subtrees' <- mapM (mapTreeM action) subtrees
       a' <- action a
       return $ Tree a' subtrees'

-- That whole business with getting the state and putting the successor
-- in as the replacement can be abstracted out.  This action is like a 
-- post-increment operator.    
postIncrement :: Enum s => State s s
postIncrement = do val <- get
                   put (succ val)
                   return val

-- Now tag can be easily written in terms of those.
tag init tree = evalState (mapTreeM step tree) init
    where step a = do tag <- postIncrement
                      return (a, tag)

You can make mapTreeM process the local value before the subtrees if you want:

mapTreeM action (Tree a subtrees) =
    do a' <- action a
       subtrees' <- mapM (mapTreeM action) subtrees
       return $ Tree a' subtrees'

And using Control.Monad you can turn this into a one-liner:

mapTreeM action (Tree a subtrees) =
    -- Apply the Tree constructor to the results of the two actions
    liftM2 Tree (action a) (mapM (mapTreeM action) subtrees)

-- in the children-first order:
mapTreeM' action (Tree a subtrees) =
    liftM2 (flip Tree) (mapM (mapTreeM action) subtrees) (action a)

Answer 2

Taking advantage of Data.Traversable and some useful GHC extensions, we can refactor sacundim's solution further:

{-# LANGUAGE DeriveFunctor, DeriveFoldable, DeriveTraversable #-}

import Control.Monad.State
import Data.Foldable
import Data.Traversable

data Tree a = Tree a [Tree a]
  deriving (Show, Functor, Foldable, Traversable)

postIncrement :: Enum s => State s s
postIncrement = do val <- get
                   put (succ val)
                   return val

-- Works for any Traversable, not just trees!
tag :: (Enum s, Traversable t) => s -> t a -> t (a, s)
tag init tree = evalState (traverse step tree) init
    where step a = do tag <- postIncrement
                      return (a, tag)