Memoize a function of type () -> 'a

1.4k views Asked by At

This memoize function fails on any functions of type () -> 'a at runtime with a Null-Argument-Exception.

let memoize f =
    let cache = System.Collections.Generic.Dictionary()
    fun x ->
        if cache.ContainsKey(x) then 
            cache.[x]
        else 
            let res = f x
            cache.[x] <- res
            res

Is there a way to write a memoize function that also works for a () -> 'a ?

(My only alternative for now is using a Lazy type. calling x.Force() to get the value.)

4

There are 4 answers

7
Tomas Petricek On BEST ANSWER

The reason why the function fails is that F# represents unit () using null of type unit. The dictionary does not allow taking null values as keys and so it fails.

In your specific case, there is not much point in memoizing function of type unit -> 'a (because it is better to use lazy for this), but there are other cases where this would be an issue - for example None is also represented by null so this fails too:

let f : int option -> int = memoize (fun a -> defaultArg a 42)
f None

The easy way to fix this is to wrap the key in another data type to make sure it is never null:

type Key<'K> = K of 'K

Then you can just wrap the key with the K constructor and everything will work nicely:

let memoize f =
    let cache = System.Collections.Generic.Dictionary()
    fun x ->
        if cache.ContainsKey(K x) then 
            cache.[K x]
        else 
            let res = f x
            cache.[K x] <- res
            res
0
Goswin Rothenthal On

I just found that the last memoize function on the same website using Map instead of Dictionary works for 'a Option -> 'b and () -> 'a too:

let memoize1 f =
    let cache = ref Map.empty

    fun x ->
        match cache.Value.TryFind(x) with
        | Some res -> res
        | None ->
            let res = f x
            cache.Value <- cache.Value.Add(x, res)
            res
0
Gene Belitski On

Memoization having a pure function (not just of type unit -> 'a, but any other too) as a lookup key is impossible because functions in general do not have equality comparer for the reason.

It may seem that for this specific type of function unit -> 'a it would be possible coming up with a custom equality comparer. But the only approach for implementing such comparer beyond extremes (reflection, IL, etc.) would be invoking the lookup function as f1 = f2 iff f1() = f2(), which apparently nullifies any performance improvement expected from memoization.

So, perhaps, as it was already noted, for this case optimizations should be built around lazy pattern, but not memoization one.

UPDATE: Indeed, after second look at the question all talking above about functions missing equality comparer is correct, but not applicable, because memoization happens within each function's individual cache from the closure. On the other side, for this specific kind of functions with signature unit->'a, i.e. at most single value of argument, using Dictionary with most one entry is an overkill. The following similarly stateful, but simpler implementation with just one memoized value will do:

let memoize2 f =
    let notFilled = ref true
    let cache = ref Unchecked.defaultof<'a>
    fun () ->
        if !notFilled then
            cache := f ()
            notFilled := false
        !cache

used as let foo = memoize2(fun () -> ...heavy on time and/or space calculation...) with first use foo() performing and storing the result of calculation and all successive foo() just reusing the stored value.

1
Dzmitry Lahoda On

Solution with mutable dictionary and single dictionary lookup call: let memoize1 f = // printfn "Dictionary" let cache = System.Collections.Generic.Dictionary() fun x -> let result, value = cache.TryGetValue(x) match result with | true -> value | false -> // printfn "f x" let res = f x cache.Add(x, res) res