Haskell idiom for stateful tree transformation

Question

data Tree a b = Leaf a | Branch b [Tree a b]

Given a pair of functions f :: a -> a' and g :: b -> b' I can easily transform a Tree a b into a Tree a' b'.

type Transform a b = a -> b
treeTransform :: Transform leaf leaf' ->
                 Transform branch branch' ->
                 Tree leaf branch ->
                 Tree leaf' branch'
treeTransform f _ (Leaf a) = Leaf (f a)
treeTransform f g (Branch b ts) = Branch (g b) (map (treeTransform f g) ts)

This tree is a bifunctor, and treeTransform above is just a bimap. Nothing special.

Now what happens when I need to thread state through f and g?

type StatefulTransform s a a' = s -> a -> (s, a')

statefulTreetransform :: StatefulTransform state leaf leaf' ->
                         StatefulTransform state branch branch' ->
                         state ->
                         Tree leaf branch ->
                         Tree leaf' branch'

Now there's more than one way to implement this function because there are different ways to traverse the tree.

I can implement the transformation using depth-first traversal, but breadth-first traversal is a stumbling block. Extracting data from the tree into a list breadth-first is relatively easy. Transforming extracted data is also straightforward. But how do I bend the transformed data back into the original tree shape?

Answer 1

There are multiple traversal orders even if you're not trying to do the bifunctor thing! I'll discuss how you might do this for a boring old functor; the extra type parameter can be handled, too, but doing so distracts from the core idea. So here's my boring-functor tree type:

data Tree x = Node x [Tree x]

The traditional way to do a breadth first traversal is to, as an intermediate step, produce a list of lists. There's one element per level of the tree for the outer list. Like this:

notQuiteBF :: Tree x -> [[x]]
notQuiteBF (Node x children) = [x] : (map concat . transpose . map notQuiteBF) children

Then the actual breadth-first traversal is just the concatenation of these lists.

bf :: Tree x -> [x]
bf = concat . notQuiteBF

The thing that's nice about [x] is it's enough information to iterate over the values in the tree. Unfortunately, it's not quite enough information to know how to reorder the traversals from multiple children: we know a breadth-first ordering of the first child's nodes and of the second child's nodes, but we don't know what depth each element is, so we can't weave them together.

Some clever fellow asked this question: what if we just remembered that depth inforation? So in notQuiteBF, we use a richer structure. The thing that's nice about [[x]] is it's enough information to reorder the elements, even though we constructed it from an essentially depth-first visitation of the tree's nodes. Unfortunately, it's not quite enough information to reconstruct the shape of the tree if we needed to: we know what the sequence of elements is at each level, but we don't know which parent each of those elements is associated with.

So now I ask: what if we just remember that extra info? Here's how: instead of [[x]], we'll return [[[x]]] as our intermediate structure. The outer list has one element per depth, as before. The next layer has one element per node in the previous depth; and the final layer has the children associated with that parent.

Let's see an example:

                       a
                      / \
                     /   \
                    /     \
                   b       c
                  / \      |
                 d   e     f
                     |    / \
                     g   h   i

For this tree, we get the following list of lists of lists, with suggestive whitespace:

[[[a                    ]]
,[ [b        ,c        ]]
,[  [d ,e   ],[f      ]]
,[   [],[g ],  [h ,i ]]
,[       [],    [],[]]
]

Wellllll... to build the tree back up, we'd actually sort of prefer to have that in the reverse order.

[[[],[],[]]
,[[],[g],[h,i]]
,[[d,e],[f]]
,[[b,c]]
,[[a]]
]

Let's write the reconstruction algorithm first.

rebuild :: [[[x]]] -> [Tree x]
rebuild = concat . go [] where
    go trees [] = trees
    go trees (xss:xsss) = go (weirdZipWith Node xss trees) xsss

weirdZipWith :: (x -> y -> z) -> [[x]] -> [y] -> [[z]]
weirdZipWith f [] _ = []
weirdZipWith f ([]:xss) ys = [] : weirdZipWith f xss ys
weirdZipWith f _ [] = []
weirdZipWith f ((x:xs):xss) (y:ys)
    = let (b, e) = splitAt 1 (weirdZipWith f (xs:xss) ys)
      in map (f x y:) b ++ e

Try it out in ghci:

> rebuild [["","",""],["","g","hi"],["de","f"],["bc"],["a"]]
[Node 'a' [Node 'b' [Node 'd' [],Node 'e' [Node 'g' []]],Node 'c' [Node 'f' [Node 'h' [],Node 'i' []]]]]

Looks good. Now the other direction. It's a very slight variation on notQuiteBF above.

bf :: Tree x -> [[[x]]]
bf (Node x children) = [[x]] : [concat (concat b)] : e where
    (b, e) = splitAt 1 . map concat . transpose . map bf $ children

We could double-check our work:

> quickCheck (\t -> (rebuild . reverse . bf) t == [t :: Tree Int])
+++ OK, passed 100 tests.

With those tools in place, it's pretty easy to write an Applicative traversal: we'll just build up the list of elements in the correct order, call f on each of them while preserving the list structure, then rebuild the tree. So:

bfTraverse :: Applicative f => (x -> f y) -> Tree x -> f (Tree y)
bfTraverse f = id
    . fmap (head . rebuild . reverse)
    . traverse (traverse (traverse f))
    . bf

(It might take quite a subtle argument to become convinced that head is safe here!) Try it out in ghci:

> bfTraverse (\x -> putStrLn [x] >> pure (toUpper x)) (Node 'a' [Node 'b' [Node 'd' [],Node 'e' [Node 'g' []]],Node 'c' [Node 'f' [Node 'h' [],Node 'i' []]]])
a
b
c
d
e
f
g
h
i
Node 'A' [Node 'B' [Node 'D' [],Node 'E' [Node 'G' []]],Node 'C' [Node 'F' [Node 'H' [],Node 'I' []]]]