## 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