Gather data about existing tree-like data

Question

Let's say we have existing tree-like data and we would like to add information about depth of each node. How can we easily achieve that?

Data Tree = Node Tree Tree | Leaf

For each node we would like to know in constant complexity how deep it is. We have the data from external module, so we have information as it is shown above. Real-life example would be external HTML parser which just provides the XML tree and we would like to gather data e.g. how many hyperlinks every node contains.

Functional languages are created for traversing trees and gathering data, there should be an easy solution.

Obvious solution would be creating parallel structure. Can we do better?

Answer 1

The standard trick, which I learned from Chris Okasaki's wonderful Purely Functional Data Structures is to cache the results of expensive operations at each node. (Perhaps this trick was known before Okasaki's thesis; I don't know.) You can provide smart constructors to manage this information for you so that constructing the tree need not be painful. For example, when the expensive operation is depth, you might write:

module SizedTree (SizedTree, sizedTree, node, leaf, depth) where

data SizedTree = Node !Int SizedTree SizedTree | Leaf

node l r = Node (max (depth l) (depth r) + 1) l r
leaf = Leaf

depth (Node d _ _) = d
depth Leaf = 0

-- since we don't expose the constructors, we should
-- provide a replacement for pattern matching
sizedTree f v (Node _ l r) = f l r
sizedTree f v Leaf = v

Constructing SizedTrees costs O(1) extra work at each node (hence it is O(n) work to convert an n-node Tree to a SizedTree), but the payoff is that checking the depth of a SizedTree -- or of any subtree -- is an O(1) operation.

Answer 2

You do need some another data where you can store these Ints. Define Tree as

data Tree a = Node Tree a Tree | Leaf a

and then write a function

annDepth :: Tree a -> Tree (Int, a)

Your original Tree is Tree () and with pattern synonyms you can recover nice constructors.

---

If you want to preserve the original tree for some reason, you can define a view:

{-# LANGUAGE GADTs, DataKinds #-}

data Shape = SNode Shape Shape | SLeaf

data Tree a sh where
    Leaf :: a -> Tree a SLeaf
    Node :: Tree a lsh -> a -> Tree a rsh -> Tree a (SNode lsh rsh)

With this you have a guarantee that an annotated tree has the same shape as the unannotated. But this doesn't work good without proper dependent types.

---

Also, have a look at the question Boilerplate-free annotation of ASTs in Haskell?

Answer 3

The standard solution is what @DanielWagner suggested, just extend the data structure. This can be somewhat inconvenient, but can be solved: Smart constructors for creating instances and using records for pattern matching.

Perhaps Data types a la carte could help, although I haven't used this approach myself. There is a library compdata based on that.

A completely different approach would be to efficiently memoize the values you need. I was trying to solve a similar problem and one of the solutions is provided by the library stable-memo. Note that this isn't a purely functional approach, as the library is internally based on object identity, but the interface is pure and works perfectly for the purpose.