## Memoizing polymorphic functions – part two

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?

## Neel Krishnaswami:

Hi Conal, I don’t know if this is enough for your purpose, but there’s a beautiful trick I learned from Janis Voigtlander’s 2009 POPL paper “Bidirectionalization for Free”. The idea is to to exploit parametricity — if you have a function that operates over lists parametrically, you can find out what it does on

anyargument by passing it a list of unique integers, and looking at how it permutes those integers.The same trick ought to work for memoization. If we want to write a function

then it’s sufficient if we can write the memo function for the Int case. More abstactly, we can do this if k and v are functors, and k also has a way to distribute over the state monad (that is, we want a function

`seq : (k (state a)) -> (state (k a))`

. Then, we want to write something like (I’ll use a mutant hybrid ML/Haskell syntax here):The idea is that we use map and seq to construct an action that will walk over the key, replacing each polymorphic value with a unique integer index. Map crawls over the term, and we stick a little command in a state monad to eat the value there. Then, we use

`seq`

to sequence all those terms, ensuring we a) get unique indices at all the leaves, and b) we can get a mapping from those indices back to the original values.Then, we can use memoize the int version of the value, and its result. Given an integer-valued result, we can use the mapping from indices to values to put back the correct values. So v needs to be a functor, and k needs to be a functor with a sequencing distribution property — the traversable stuff ought to be the right general property to ask for.

18 June 2009, 11:30 am## Conal Elliott » Blog Archive » Memoizing polymorphic functions via unmemoization:

[...] combination of patricia trees (for integer-keyed maps), stable names, and type equality proofs. The second post showed how to transform some functions that do not quite fit the required form so that they do [...]

2 October 2010, 7:41 am