Numbering nodes in a tree with and without lenses.

TreeNumbering.hs

module TreeNumbering where

import Control.Lens
import Control.Monad.State
import Data.Tree

tree_ :: Tree Char
tree_ = Node 'd'
  [ Node 'b'
    [ Node 'a' []
    , Node 'c' []
    ]
  , Node 'f'
    [ Node 'e' []
    , Node 'g' []
    ]
  ]

tag :: Traversable t => t a -> t (Int, a)
tag t = evalState (traverse step t) 0
  where
    step a = do
      tag <- postIncrement
      pure (tag, a)

    postIncrement = do
      result <- get
      modify succ
      pure result

lensTag :: Tree a -> Tree (Int, a)
lensTag = unsafePartsOf each %~ zip [0..]

--  λ> putStrLn . drawTree . fmap show . tag $ tree_
--  (0,'d')
--  |
--  +- (1,'b')
--  |  |
--  |  +- (2,'a')
--  |  |
--  |  `- (3,'c')
--  |
--  `- (4,'f')
--     |
--     +- (5,'e')
--     |
--     `- (6,'g')

--  λ> putStrLn . drawTree . fmap show . lensTag $ tree_
--  (0,'d')
--  |
--  +- (1,'b')
--  |  |
--  |  +- (2,'a')
--  |  |
--  |  `- (3,'c')
--  |
--  `- (4,'f')
--     |
--     +- (5,'e')
--     |
--     `- (6,'g')

--  λ> tag tree_ == lensTag tree_
--  True