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

1. #### 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 any argument 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

```     memo' : (all a. k a -> v a) -> (all a. k a -> v a)`
```

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):

```type state r a = (Int, Int -> r) -> ((Int, Int -> r), a)

eat : r -> state r int
eat r (size, mapping) = ((size+1, \n. if n = size then r else mapping n), size)

run : state r a -> int -> r -> ((int, int -> r), a)
run state initsize initmap = state (initsize, initmap)

seq : k (state r a) -> state r (k a)
seq blah = whatever

memo' (f : all a. k a -> v a) =
let table : ref (k int -> option (v int)) =  ref (\x. Nothing)   -- allocate a table
let memofun x : k b =
let leaves : k (state b int) = map_k eat x
let traversal : state b (k int) = seq leaves
let ((size, indices), shape) = run traversal 0 (\x. bot)
case !table shape of
Just result_shape -> fmap indices result_shape
| Nothing -> let result_shape = f shape in
begin
table := update table shape result_shape;  -- imperatively update the memo table
return (fmap indices result_shape)
end

```

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.

2. #### 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 [...]