Weak-memoizing result of multi-parameter function in OCaml

Question

I am looking for a way to memoize the results of an OCaml function f that takes two parameters (or more, in general). In addition (and this is the difficult part), I want the map underlying this process to forget a result altogether if either of the values for the two parameters is garbage collected.

For a function that takes exactly one argument, this can be done with the Weak module and its Make functor in a straightforward way. To generalise this to something that can memoize functions of higher arity, a naive solution is to create a weak map from tuples of values to result values. But this will not work correctly with respect to garbage collection, as the tuple of values only exists within the scope of the memoization function, not the client code that calls f. In fact, the weak reference will be to the tuple, which is going to be garbage collected right after memoization (in the worst case).

Is there a way to do this without re-implementing Weak.Make?

Hash-consing is orthogonal to my requirements and is, in fact, not really desirable for my values.

Thanks!

Answer 1

Instead of indexing by tuples you could have a tree structure. You'd have one weak table indexed by the first function parameter whose entries are secondary weak tables. The secondary tables would be indexed by the second function parameter and contain the memoized results. This structure will forget the memoized function results as soon as either function parameter is GCed. However, the secondary tables themselves will be retained as long as the first function parameter is live. Depending on the sizes of your function results and the distribution of different first parameters, this could be a reasonable tradeoff.

I haven't tested this, either. Also it seems reasonably obvious.

Answer 2

One idea is to perform your own garbage collection.

For simplicity, let's assume that all arguments have the same type k.

In addition to the main weak table containing the memoized results keyed by k * k, create a secondary weak table containing single arguments of type k. The idea is to scan the main table once in a while and to remove the bindings that are no longer wanted. This is done by looking up the arguments in the secondary table; then if any of them is gone you remove the binding from the main table.

(Disclaimer: I haven't tested this; it may not work or there may be better solutions)