lenses_on_custom_trees.hs

{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE RankNTypes      #-}

{-  mutate = set (on 0 . on 1 . on 1 . item) (* 999)


    1                                 1
    |                                 |
    +- 2                              +- 2
    |  |                              |  |
    |  +- 4                           |  +- 4
    |  |  |                           |  |  |
    |  |  +- 8                        |  |  +- 8
    |  |  |                           |  |  |
    |  |  `- 12                       |  |  `- 12
    |  |                              |  |
    |  `- 6                           |  `- 6
    |     |                           |     |
    |     +- 12                       |     +- 12
    |     |              `mutate`     |     |
    |     `- 17982     <<<<=======    |     `- 18
    |                                 |
    `- 3                              `- 3
       |                                 |
       +- 6                              +- 6
       |  |                              |  |
       |  +- 12                          |  +- 12
       |  |                              |  |
       |  `- 18                          |  `- 18
       |                                 |
       `- 9                              `- 9
          |                                 |
          +- 18                             +- 18
          |                                 |
          `- 27                             `- 27
-}

import Control.Lens
import Control.Applicative
import Data.Tree
import Data.Tree.Lens

myTree :: Tree Integer
myTree = unfoldTree (\x -> (x, [x*2, x*3])) 1

chop :: Int -> Tree a -> Tree a
chop 0 (Node x _) = Node x []
chop n (Node x f) = Node x (map (chop (pred n)) f)

on :: Applicative f => Int -> (Tree a -> f (Tree a)) -> Tree a -> f (Tree a)
on n = branches . ix n

mutate :: Tree Integer -> Tree Integer
mutate = over (on 0 . on 1 . on 1 . root) (* 999)

main :: IO ()
main = putStrLn $ drawTree $ fmap show $ chop 3 $ mutate $ myTree