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.