Elegant memoization with higher-order types

A while back, I got interested in functional memoization, especially after seeing some code from Spencer Janssen using the essential idea of Ralf Hinze’s paper Generalizing Generalized Tries. The blog post Elegant memoization with functional memo tries describes a library, MemoTrie, based on both of these sources, and using associated data types. I would have rather used associated type synonyms and standard types, but I couldn’t see how to get the details to work out. Recently, while playing with functor combinators, I realized that they might work for memoization, which they do quite nicely.

This blog post shows how functor combinators lead to an even more elegant formulation of functional memoization. The code is available as part of the functor-combo package.

The techniques in this post are not so much new as they are ones that have recently been sinking in for me. See Generalizing Generalized Tries, as well as Generic programming with fixed points for mutually recursive datatypes.

Edits:

  • 2011-01-28: Fixed small typo: “b^^a^^” ⟼ “ba
  • 2010-09-10: Corrected Const definition to use newtype instead of data.
  • 2010-09-10: Added missing Unit type definition (as Const ()).

Tries as associated data type

The MemoTrie library is centered on a class HasTrie with an associated data type of tries (efficient indexing structures for memoized functions):

class HasTrie k where
    data (:→:) k :: * → *
    trie   :: (k  →  v) → (k :→: v)
    untrie :: (k :→: v) → (k  →  v)

The type a :→: b represents a trie that maps values of type a to values of type b. The trie representation depends only on a.

Memoization is a simple combination of these two methods:

memo :: HasTrie a ⇒ (a → b) → (a → b)
memo = untrie . trie

The HasTrie instance definitions correspond to isomorphisms invoving function types. The isomorphisms correspond to the familiar rules of exponents, if we translate a → b into ba. (See Elegant memoization with functional memo tries for more explanation.)

instance HasTrie () where
    data () :→: x = UnitTrie x
    trie f = UnitTrie (f ())
    untrie (UnitTrie x) = const x

instance (HasTrie a, HasTrie b) ⇒ HasTrie (Either a b) where
    data (Either a b) :→: x = EitherTrie (a :→: x) (b :→: x)
    trie f = EitherTrie (trie (f . Left)) (trie (f . Right))
    untrie (EitherTrie s t) = either (untrie s) (untrie t)

instance (HasTrie a, HasTrie b) ⇒ HasTrie (a,b) where
    data (a,b) :→: x = PairTrie (a :→: (b :→: x))
    trie f = PairTrie (trie (trie . curry f))
    untrie (PairTrie t) = uncurry (untrie .  untrie t)

Functors and functor combinators

For notational convenience, let “(:→:)” be a synonym for “Trie“:

type k :→: v = Trie k v

And replace the associated data with an associated type.

class HasTrie k where
    type Trie k :: * → *
    trie   :: (k  →  v) → (k :→: v)
    untrie :: (k :→: v) → (k  →  v)

Then, imitating the three HasTrie instances above,

type Trie () v = v

type Trie (Either a b) v = (Trie a v, Trie b v)

type Trie (a,b) v = Trie a (Trie b v)

Imagine that we have type lambdas for writing higher-kinded types.

type Trie () = λ v → v

type Trie (Either a b) = λ v → (Trie a v, Trie b v)

type Trie (a,b) = λ v → Trie a (Trie b v)

Type lambdas are often written as “Λ” (capital “λ”) instead. In the land of values, these three right-hand sides correspond to common building blocks for functions, namely identity, product, and composition:

id      = λ v → v
f *** g = λ v → (f v, g v)
g  .  f = λ v → g (f v)

These building blocks arise in the land of types.

newtype Id a = Id a

data (f :*: g) a = f a :*: g a

newtype (g :. f) a = O (g (f a))

where Id, f and g are functors. Sum and a constant functor are also common building blocks:

data (f :+: g) a = InL (f a) | InR (g a)

newtype Const x a = Const x

type Unit = Const () -- one non-⊥ inhabitant

Tries as associated type synonym

Given these standard definitions, we can eliminate the special-purpose data types used, replacing them with our standard functor combinators:

instance HasTrie () where
  type Trie ()  = Id
  trie   f      = Id (f ())
  untrie (Id v) = const v

instance (HasTrie a, HasTrie b) => HasTrie (Either a b) where
  type Trie (Either a b) = Trie a :*: Trie b
  trie   f           = trie (f . Left) :*: trie (f . Right)
  untrie (ta :*: tb) = untrie ta `either` untrie tb

instance (HasTrie a, HasTrie b) ⇒ HasTrie (a , b) where
  type Trie (a , b) = Trie a :. Trie b
  trie   f      = O (trie (trie . curry f))
  untrie (O tt) = uncurry (untrie . untrie tt)

At first blush, it might appear that we’ve simply moved the data type definitions outside of the instances. However, the extracted functor combinators have other uses, as explored in polytypic programming. I’ll point out some of these uses in the next few blog posts.

Isomorphisms

Many types are isomorphic variations, and so their corresponding tries can share a common representation. For instance, triples are isomorphic to nested pairs:

detrip :: (a,b,c) → ((a,b),c)
detrip (a,b,c) = ((a,b),c)

trip :: ((a,b),c) → (a,b,c)
trip ((a,b),c) = (a,b,c)

A trie for triples can be a a trie for pairs (already defined). The trie and untrie methods then just perform conversions around the corresponding methods on pairs:

instance (HasTrie a, HasTrie b, HasTrie c) ⇒ HasTrie (a,b,c) where
    type Trie (a,b,c) = Trie ((a,b),c)
    trie f = trie (f . trip)
    untrie t = untrie t . detrip

All type isomorphisms can use this same pattern. I don’t think Haskell is sufficiently expressive to capture this pattern within the language, so I’ll resort to a C macro. There are five parameters:

  • Context: the instance context;
  • Type: the type whose instance is being defined;
  • IsoType: the isomorphic type;
  • toIso: conversion function to IsoType; and
  • fromIso: conversion function from IsoType.

The macro:

#define HasTrieIsomorph(Context,Type,IsoType,toIso,fromIso)  
instance Context ⇒ HasTrie (Type) where {  
  type Trie (Type) = Trie (IsoType);  
  trie f = trie (f . (fromIso));  
  untrie t = untrie t . (toIso);  
}

Now we can easily define HasTrie instances:

HasTrieIsomorph( (), Bool, Either () ()
               ,  c -> if c then Left () else Right ()
               , either ( () -> True) ( () -> False))

HasTrieIsomorph( (HasTrie a, HasTrie b, HasTrie c), (a,b,c), ((a,b),c)
               , λ (a,b,c) → ((a,b),c), λ ((a,b),c) → (a,b,c))

HasTrieIsomorph( (HasTrie a, HasTrie b, HasTrie c, HasTrie d)
               , (a,b,c,d), ((a,b,c),d)
               , λ (a,b,c,d) → ((a,b,c),d), λ ((a,b,c),d) → (a,b,c,d))

In most (but not all) cases, the first argument (Context) could simply be that the isomorphic type HasTrie, e.g.,

HasTrieIsomorph( HasTrie ((a,b),c), (a,b,c), ((a,b),c)
               , λ (a,b,c) → ((a,b),c), λ ((a,b),c) → (a,b,c))

We could define another macro that captures this pattern and requires one fewer argument. On the other hand, there is merit to keeping the contextual requirements explicit.

Regular data types

A regular data type is one in which the recursive uses are at the same type. Functions over such types are often defined via monomorphic recursion. Data types that do not satisfy this constraint are called “nested“.

As in several recent generic programming systems, regular data types can be encoded generically through a type class that unwraps one level of functor from a type. The regular data type is the fixpoint of that functor. See, e.g., Polytypic programming in Haskell. Adopting the style of A Lightweight Approach to Datatype-Generic Rewriting,

class Functor (PF t) ⇒ Regular t where
  type PF t :: * → *
  wrap   :: PF t t → t
  unwrap :: t → PF t t

Here “PF” stands for “pattern functor”.

The pattern functors can be constructed out of the functor combinators above. For instance, a list at the top level is either empty or a value and a list. Translating this description:

instance Regular [a] where
  type PF [a] = Unit :+: Const a :*: Id

  unwrap []     = InL (Const ())
  unwrap (a:as) = InR (Const a :*: Id as)

  wrap (InL (Const ()))          = []
  wrap (InR (Const a :*: Id as)) = a:as

As another example, consider rose trees ([Data.Tree][]):

data Tree  a = Node a [Tree a]

instance Regular (Tree a) where

  type PF (Tree a) = Const a :*: []

  unwrap (Node a ts) = Const a :*: ts

  wrap (Const a :*: ts) = Node a ts

Regular types allow for even more succinct HasTrie instance implementations. Specialize HasTrieIsomorph further:

#define HasTrieRegular(Context,Type)  
HasTrieIsomorph(Context, Type, PF (Type) (Type) , unwrap, wrap)

For instance, for lists and rose trees:

HasTrieRegular(HasTrie a, [a])
HasTrieRegular(HasTrie a, Tree a)

The HasTrieRegular macro could be specialized even further for single-parameter polymorphic data types:

#define HasTrieRegular1(TypeCon) HasTrieRegular(HasTrie a, TypeCon a)

HasTrieRegular1([])
HasTrieRegular1(Tree)

You might wonder if I’m cheating here, by claiming very simple trie specifications when I’m really just shuffling code around. After all, the complexity removed from HasTrie instances shows up in Regular instances. The win in making this shuffle is that Regular is handy for other purposes, as illustrated in Generic programming with fixed points for mutually recursive datatypes (including fold, unfold, and fmap). (More examples in A Lightweight Approach to Datatype-Generic Rewriting.)

Trouble

Sadly, these elegant trie definitions have a problem. Trying to compile them leads to a error message from GHC. For instance,

Nested type family application
  in the type family application: Trie (PF [a] [a])
(Use -XUndecidableInstances to permit this)

Adding UndecidableInstances silences this error message, but leads to nontermination in the compiler.

Expanding definitions, I can see the likely cause of nontermination. The definition in terms of a type family allows an infinite type to sneak through, and I guess GHC’s type checker is unfolding infinitely.

As a simpler example:

{-# LANGUAGE TypeFamilies, UndecidableInstances #-}

type family List a :: *

type instance List a = Either () (a, List a)

-- Hangs ghc 6.12.1:
nil :: List a
nil = Left ()

A solution

Since GHC’s type-checker cannot handle directly recursive types, perhaps we can use a standard avoidance strategy, namely introducing a newtype or data definition to break the cycle. For instance, as a trie for [a], we got into trouble by using the trie of the unwrapped form of [a], i.e., Trie (PF [a] [a]). So instead,

newtype ListTrie a v = ListTrie (Trie (PF [a] [a]) v)

which is to say

newtype ListTrie a v = ListTrie (PF [a] [a] :→: v)

Now wrap and unwrap as before, and add & remove ListTrie as needed:

instance HasTrie a ⇒ HasTrie [a] where
  type Trie [a] = ListTrie a
  trie f = ListTrie (trie (f . wrap))
  untrie (ListTrie t) = untrie t . unwrap

Again, abstract the boilerplate code into a C macro:

#define HasTrieRegular(Context,Type,TrieType,TrieCon) 
newtype TrieType v = TrieCon (PF (Type) (Type) :→: v); 
instance Context ⇒ HasTrie (Type) where { 
  type Trie (Type) = TrieType; 
  trie f = TrieCon (trie (f . wrap)); 
  untrie (TrieCon t) = untrie t . unwrap; 
}

For instance,

HasTrieRegular(HasTrie a, [a] , ListTrie a, ListTrie)
HasTrieRegular(HasTrie a, Tree, TreeTrie a, TreeTrie)

Again, simplify a bit with a specialization to unary regular types:

#define HasTrieRegular1(TypeCon,TrieCon) 
HasTrieRegular(HasTrie a, TypeCon a, TrieCon a, TrieCon)

And then use the following declarations instead:

HasTrieRegular1([]  , ListTrie)
HasTrieRegular1(Tree, TreeTrie)

Similarly for binary etc as needed.

The second macro parameter (TrieCon) is just a name, which I don’t to be used other than in the macro-generated code. It could be eliminated, if there were a way to gensym the name. Perhaps with Template Haskell?

Conclusion

I like the elegance of constructing memo tries in terms of common functor combinators. Standard pattern functors allow for extremely succinct trie specifications for regular data types. However, these specifications lead to nontermination of the type checker, which can then be avoided by the standard trick of introducing a newtype to break type recursion. As often, this trick brings introduces some clumsiness. Perhaps the problem can also be avoided by using a formulation using bifunctors, as in Design Patterns as Higher-Order Datatype-Generic Programs and Polytypic programming in Haskell, which allows the fixed-point nature of regular data types to be exposed.

6 Comments

  1. Cain Norris:

    This certainly seems to be amenable to Template Haskell, since TH provides fresh name generation and can handle the simple CPP stuff. I’ve been looking for a little project to brush up on TH, this could be a fun start :-)

  2. Conal Elliott » Blog Archive » Memoizing higher-order functions:

    […] About « Elegant memoization with higher-order types […]

  3. Conal Elliott » Blog Archive » Another angle on zippers:

    […] in Elegant memoization with higher-order types, let’s use the following […]

  4. Conal Elliott » Blog Archive » Differentiation of higher-order types:

    […] use the same set of functor combinators as in Elegant memoization with higher-order types and Memoizing higher-order […]

  5. Conal Elliott » Blog Archive » Memoizing polymorphic functions via unmemoization:

    […] or, using functor combinators, […]

  6. Conal Elliott » Blog Archive » Composable parallel scanning:

    […] To see how to scan over a broad range of functors, let’s look at each of the functor combinators, e.g., as in Elegant memoization with higher-order types. […]

Leave a comment