Labeling traversables.

Label.hs

{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE DerivingStrategies #-}

import Control.Monad.State

data Tree a = Node (Tree a) a (Tree a) | Leaf
  deriving stock (Show, Functor, Foldable, Traversable)

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

label :: Tree a -> Tree (a, Int)
label = fst . go 0
  where
    go n = \case
      Leaf -> (Leaf, n)
      Node lhs val rhs ->
        let (lhs', o) = go n lhs
            (val', p) = ((val, o), o + 1)
            (rhs', q) = go p rhs
         in (Node lhs' val' rhs', q)

label' :: Tree a -> Tree (a, Int)
label' t = evalState (traverse f t) 0
  where
    f x = do
      n <- get
      put (n + 1)
      pure (x, n)

labelTraversable :: (Traversable t) => t a -> t (a, Int)
labelTraversable t = evalState (traverse f t) 0
  where
    f x = do
      n <- get
      put (n + 1)
      pure (x, n)

foo :: Tree Char
foo = Node Leaf 'a' (Node Leaf 'b' Leaf)

bar :: Rose Char
bar = Rose 'a' [Rose 'b' []]

main :: IO ()
main = do
  print (label foo)
  print (label' foo)
  print (labelTraversable foo)
  print (labelTraversable bar)