Part one of this series introduced the problem of memoizing functions involving polymorphic recursion. The caching data structures used in memoization typically handle only one type of argument at a time. For instance, one can have finite maps of differing types, but each concrete finite map holds just one type of key and one type of value.

I extended memoization to handle polymorphic recursion by using an existential type together with a reified type of types. This extension works (afaik), but it is restricted to a particular form for the type of the polymorphic function being memoized, namely

-- Polymorphic function
type k :--> v = forall a. HasType a => k a -> v a

My motivating example is a GADT-based representation of typed lambda calculus, and some of the functions I want to memoize do not fit the pattern. After writing part one, I fooled around and found that I could transform these awkwardly typed polymorphic functions into isomorphic form that does indeed fit the restricted pattern of polymorphic types I can handle.

Awkward types

The first awkwardly typed memoizee is the function application constructor:

type AT = forall a b . (HasType a, HasType b) => E (a -> b) -> E a -> E b

(:^) :: AT

Right away AT misses the required form. It has two HasType constraints, and the first argument is parameterized over two type variables instead of one. However, the second argument looks more promising, so let’s flip the arguments to get an isomorphic type:

forall a b . (HasType a, HasType b) => E a -> E (a -> b) -> E b

And then move the quantifier and constraint on b inside the outer (first) ->:

forall a . HasType a => E a -> (forall b. HasType b => E (a -> b) -> E b)

We’re getting closer. Next, define a newtype wrapper.

newtype A2 a = A2 (forall b. HasType b => E (a -> b) -> E b)

So that AT is isomorphic to

forall a . HasType a => E a -> A2 a

i.e.,

E :--> A2

The function inside of A2 doesn’t have the required form, but another newtype wrapper finishes the job.

newtype EF a b = EF {unEF :: E (a -> b) }

type H' a = EF a :--> E

newtype H a = H { unH :: H' a }

The AT type is isomorphic AP where

type AP = E :--> H

Curried memoization

A “curried memo function” is one that takes one argument and produces another memo function. For a simple memo function, not involving polymorphic recursion, there’s a simple recipe for curried memoization:

memo2 :: (a -> b -> c) -> (a -> b -> c)
memo2 f = memo (memo . f)

Our more polymorphic memo makes currying a little more awkward. First, here’s a helper function for working inside of the representation of an H:

inH :: (H' a -> H' a) -> (H a -> H a)
inH h z = H (h (unH z))

The following more elegant definition doesn’t type-check, due to the rank 2 polymorphism:

inH f = H . f . unH  -- type error

Now our AP memoizer is much like memo2:

memoAP :: AP -> AP
memoAP app' = memo (inH memo . app')

(A more general, consistent type for memoAP is (f :--> H) -> (f :--> H).)

Isomorphisms

Now, to define the isomorphisms. Define

toAP   :: AT -> AP
fromAP :: AP -> AT

The definitions:

toAP app ea = H $  (EF eab) -> app eab ea
fromAP app' eab ea = unH (app' ea) (EF eab)

If you erase the newtype wrappers & unwrappers, you’ll see that toAP and fromAP are both just flip.

I constructed fromAP from the following specification:

toAP (fromAP app') == app'

Transforming step-by-step into equivalent specifications:

 ea -> H $  (EF eab) -> (fromAP app') eab ea == app'

H $  (EF eab) -> (fromAP app') eab ea == app' ea

 (EF eab) -> (fromAP app') eab ea == unH (app' ea)

(fromAP app') eab ea == unH (app' ea) (EF eab)

fromAP app' eab ea == unH (app' ea) (EF eab)

Memoizing vis isomorphisms

Finally, I can memoize

memoAT :: AT -> AT
memoAT app = fromAP (memoAP (toAP app))

Again, a more elegant definition via (.) fails to type-check, due to rank 2 polymorphism.

The Lam (lambda abstraction) constructor can be handled similarly:

Lam  :: (HasType a, HasType b) => V a -> E b -> E (a -> b)

This time, no flip is needed.

I wonder

How far does this isomorphism trick go?

Is there an easier way to memoize polymorphic functions?

Memoization takes a function and gives back a semantically equivalent function that reuses rather than recomputes when applied to the same argument more than once. Variations include not-quite-equivalence due to added strictness, and replacing value equality with pointer equality.

Memoization is often packaged up polymorphically:

memo :: (???) => (k -> v) -> (k -> v)

For pointer-based (“lazy”) memoization, the type constraint (“???”) is empty. For equality-based memoization, we’d need at least Eq k, and probably Ord k or HasTrie k for efficient lookup (in a finite map or a possibly infinite memo trie).

Although memo is polymorphic, its argument is a monomorphic function. Implementations that use maps or tries exploit that monomorphism in that they use a type like Map k v or Trie k v. Each map or trie is built around a particular (monomorphic) type of keys. That is, a single map or trie does not mix keys of different types.

Now I find myself wanting to memoize polymorphic functions, and I don’t know how to do it.

Flavors of polymorphism

If a recursively defined function f is polymorphic, then the recursion may still be monomorphic. That is, recursive calls may be restricted to the same type instance as the parent call. Most recursive polymorphic functions fit this form, because most polymorphic recursive data types are “regular”, meaning that a polymorphic data type is included in itself only at the same type instance. For instance, the usual polymorphic lists and trees are regular.

Example: GADTs

Among other places, non-regular, or “nested data types” arise in statically typed encodings of typed languages. For instance, here’s a GADT (generalized algebraic data type) for typed lambda calculus expressions:

-- Variables
data V a = V String (Type a)

-- Expressions
data E :: * -> * where
  Lit  :: a -> E a                      -- literal
  Var  :: V  a -> E a                   -- variable
  (:^) :: E (a -> b) -> E a -> E b      -- application
  Lam  :: V a -> E b -> E (a -> b)      -- abstraction

The Type type is sort of like TypeRep, except that it is statically typed.

data Type :: * -> * where
  Bool   :: Type Bool
  Float  :: Type Float
  ...
  (:*:)  :: Type a -> Type b -> Type (a,b)
  (:->:) :: Type a -> Type b -> Type (a->b)

These GADTs (E and Type) are both non-regular, and so recursive functions over them will involve more than one argument type.

So how can we memoize?

A first try

Let’s consider a specific case of polymorphic functions:

type k :--> v = ∀ a. HasType a => k a -> v a

HasType is to Typeable as Type is to TypeRep. The memo implementation I’m playing with relies on HasType.

The memoizer can have type

memo :: (k :--> v) -> (k :--> v)

which uses rank 2 polymorphism (because of the argument type’s ∀, which cannot be moved to the outside).

My implementation of memo is similar to the discussion in Stretching the storage manager: weak pointers and stable names in Haskell, but modernized a bit and adapted for polymorphism. It uses an IntMap of lists of pairs of stable names (akin to pointers) for arguments and values for results. The idea is to first use hashStableName to get an Int to use as an IntMap key, and then linearly traverse the resulting list of binding pairs, comparing stable keys for equality. (StableName has Eq but not Ord.) Although hashStableName can map different stable names to the same hash value, collisions are rare, so the StableBind lists rarely have more than one element and the linear search is cheap.

type SM k v = I.IntMap [StableBind k v]  -- Stable map

data StableBind k v = ∀ a. HasType a => SB (StableName (k a)) (v a)

The reason for hiding the type parameter a in StableName is so that different bindings can be for different types.

The key tricky bit is managing static typing while searching for a particular StableName a in a [StableBind]. Here’s my implementation:

blookup :: ∀ k v a. HasType a =>
           StableName (k a) -> [StableBind k v] -> Maybe (v a)
blookup stk = look
 where
   look :: [StableBind k v] -> Maybe (v a)
   look [] = Nothing
   look (SB stk' v : binds') 
     | Just Refl <- tya `tyEq` typeOf2 stk', stk == stk' = Just v
     | otherwise                                         = look binds'
   tya :: Type a
   tya = typeT

The crucial magic bit is

tyEq :: Type a -> Type b -> Maybe (a :=: b)

where a :=: b represents a proof that the types a and b are the same type. The proof type has a simple GADT representation:

data (:=:) :: * -> * -> * where Refl :: a :=: a

This simple type definition ensures that only valid type-equality proofs can exist. Well, except for ⊥. The guard’s pattern match with Just Refl will force evaluation, so that ⊥ can’t sneak by us. That match also informs the type-checker that stk and stk' are stable names of the same type in that clause, which then makes stk == stk' be well-typed, and makes Just v have the required type, i.e., Maybe (v a).

Finally, the typeOf2 function is a simple helper that peels off two type constructors and extracts a Type:

typeOf2 :: HasType a => g (f a) -> Type a

Wishing for more

The type of memo above is too restrictive for my uses. It only handles polymorphic functions of type k a -> v a, and only with the single constraint HasType a.

The reason I want lazy memoization now is that I’m compressing expressions to maximize representation sharing, as John Hughes described in Lazy Memo-functions. Once sharing is maximized, pointer-based memoization works better, because equal values are pointer-equal. To compress an expression, simply use a memoized copy function, as John suggested.

compress :: HasType a => E a -> E a
compress e = mcopy e
 where
   mcopy, copy :: HasType b => E b -> E b
   -- Memo version
   mcopy = memo copy
   -- Copier, with memo-copied components
   copy (u :^ v)  = appM (mcopy u) (mcopy v)
   copy (Lam v b) = lamM v (mcopy b)
   copy e         = e
   -- memoized constructors
   appM :: (HasType a, HasType b) => E (a -> b) -> E a -> E b
   appM = memo2 (:^)
   lamM :: (HasType a, HasType b) => V a -> E b -> E (a -> b)
   lamM = memo2 Lam

The memo2 function is defined in terms of memo, using some some newtype trickery. Its type:

memo2 :: HasType a =>
         (k a -> l a -> v a) -> (k a -> l a -> v a)

But, sigh, the (:^) and Lam constructors I’m trying to memoize do not have the required types. My higher-order-polymorphic memo functions do not have flexible enough types.

And this is where I’m stuck.

I’d appreciate your ideas and suggestions.