Traverse a rose tree until some condition is met, then modify tree

Question

I have two rose trees in Haskell of m and n nodes, respectively. I want to replace the ith node of the first tree with the jth node of the second tree.

e.g.

tree 1:

           R                                                           
  _________|____________                 
  |        |           |
  A        B           C
 / \      / \         / \
D   E    F   G       H   I

tree 2:

       r
  _____|____
  |         |       
  P         Q       
 / \      / | \      
S   T    U  V  W     

then the resulting tree of tree 1 node 7 (C) replaced with tree 2 node 4 (Q) should be (assuming indexing is pre order and starting at 0)

           R                                                           
  _________|____________                 
  |        |           |
  A        B           Q
 / \      / \        / | \
D   E    F   G      U  V  W

I have tried using a zipper, but my problem is I can't workout how to get the zipper focus to the ith element,

i.e. how can i implement a function with type:

someFunc :: Tree a -> Int -> Zipper a

that take the root of the tree and traverses it (in some order) to return the zipper focussed on the ith pre order node and the context.

from the tree library I can flatten a tree to a pre order list of values using flatten, i can change this slightly to give me a list of trees e.g.

flattenToTreeList :: Tree a -> [Tree a]
flattenToTreeList t = squish t []
  where squish (Node x ts) xs = Node x ts : foldr squish xs ts

if I could do this but with a zipper then the ith element of this list would satisfy me, but I'm lost and now going round in circles.

Answer 1

Well, I was a bit off in my comment. I misread the requirements a bit, but in a way that actually cuts the dependencies down a bit.

-- lens
import Control.Lens ((^?), contexts, cosmos, elementOf, ix)

-- containers
import Data.Tree (Tree(..))

-- adjunctions
import Control.Comonad.Representable.Store (peeks)

tree1 :: Tree Char
tree1 = Node 'R' [ Node 'A' [Node 'D' [], Node 'E' []]
                 , Node 'B' [Node 'F' [], Node 'G' []]
                 , Node 'C' [Node 'H' [], Node 'I' []]
                 ]

tree2 :: Tree Char
tree2 = Node 'R' [ Node 'P' [Node 'S' [], Node 'T' []]
                 , Node 'Q' [Node 'U' [], Node 'V' [], Node 'W' []]
                 ]

-- Replace subtree i in t1 with subtree j in t2
-- returns Nothing if either index doesn't exist
replace :: Int -> Tree a -> Int -> Tree a -> Maybe (Tree a)
replace i t1 j t2 = update =<< replacement
  where
    update u = peeks (const u) <$> contexts t1 ^? ix i
    replacement = t2 ^? elementOf cosmos j

main :: IO ()
main = print $ replace 7 tree1 4 tree2

This adds a direct dependency on the adjunctions package, but it's already a transitive dependency via lens. So it's an extra import, but no additional required packages. In exchange, it doesn't need to use tree-traversals at all.

This is a bit unlike usual lens code, in that neither cosmos nor contexts are especially common, but they're great tools for manipulating substructures of self-similar data types. And that's a perfect description of replacing subtrees.

This uses pretty conceptually heavy tools, but I think the meaning comes across pretty well.