An indexed version of Plated

IndexedPlated.hs

#! /usr/bin/env runhaskell

{-# LANGUAGE Haskell2010           #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE RankNTypes            #-}

import Control.Applicative
import Control.Lens
import Control.Monad
import Data.Tree
import Prelude

--------------------------------------------------------------------------------

class IndexedPlated i a where
  indexedPlated :: i -> IndexedTraversal' i a a

itransform
  :: IndexedPlated i a
  => (i -> a -> a)
  -> i
  -> a
  -> a
itransform = itransformOf indexedPlated

itransformOf
  :: (i -> IndexedTraversal' i a a)
  -> (i -> a -> a)
  -> i
  -> a
  -> a
itransformOf l f = go where
  go i = f i . iover (l i) go

itransformM
  :: (Monad m, IndexedPlated i a)
  => (i -> a -> m a)
  -> i
  -> a
  -> m a
itransformM = itransformMOf indexedPlated

itransformMOf
  :: (Monad m)
  => (i -> IndexedLensLike i (WrappedMonad m) a a a a)
  -> (i -> a -> m a)
  -> i
  -> a
  -> m a
itransformMOf l f = go where
  go i t = imapMOf (l i) go t >>= f i

--------------------------------------------------------------------------------

instance IndexedPlated [a] (Tree a) where
  indexedPlated p f (Node label subs) =
    Node label <$> traverse (indexed f (p <> [label])) subs

printTree :: Tree Int -> IO ()
printTree = void . itransformM display []
  where
    display :: [Int] -> Tree Int -> IO (Tree Int)
    display path t = t <$ print (path ++ [rootLabel t])

main = do
  let tree = unfoldTree go 6
      go n = (n, [n | n <- [n - 1, n - 2], n >= 0])
  printTree tree