Is there a way to generalize this TrieMap code?

Question

Below is a simple Haskell program which computes equalities on trees:

import Control.Monad
import Control.Applicative
import Data.Maybe

data Tree = Leaf | Node Tree Tree

eqTree :: Tree -> Tree -> Maybe ()
eqTree Leaf         Leaf         = return ()
eqTree (Node l1 r1) (Node l2 r2) = eqTree l1 l2 >> eqTree r1 r2
eqTree _ _ = empty

Suppose you have an association list of trees [(Tree, a)], and you'd like to find the entry for a given tree. (One can think of this as a simplified version of the type class instance lookup problem.) Naively, we would have to do O(n*s) work, where n is the number of trees, and s is the size of each tree.

We can do better if we use a trie map to represent our association list:

(>.>) = flip (.)

data TreeMap a
    = TreeMap {
        tm_leaf :: Maybe a,
        tm_node :: TreeMap (TreeMap a)
      }

lookupTreeMap :: Tree -> TreeMap a -> Maybe a
lookupTreeMap Leaf       = tm_leaf
lookupTreeMap (Node l r) = tm_node >.> lookupTreeMap l >=> lookupTreeMap r

Our lookup now only takes O(s). This algorithm is a strict generalization of the previous one, since we can test for equality by creating a singleton TreeMap () and then seeing if we get back Just (). But for practical reasons, we'd prefer not to do this, since it involves building up a TreeMap and then immediately tearing it down.

Is there a way to generalize the two pieces of code above into a new function that can operate on both Tree and TreeMap? There seems to be some similarity in how the code is structured, but it is not obvious how to abstract the differences away.

Is there a way to generalize this TrieMap code?

2 Answers