Conal Elliott » Memoizing polymorphic functions – part two (original) (raw)

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?

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