F# memoization module

memoization.fsx

module memoization

open System.Collections.Generic
open System.Collections.Concurrent

let memoization fn =
    let cache = new Dictionary<_,_>()
    printfn "counter: %A" cache.Count
    
    fun arg ->
        printfn "%A is received" arg
        match cache.TryGetValue arg with
        | true, value -> value
        | false, _ ->   
            let value = fn arg
            cache.Add(arg, value)
            value

let thread_safe_memoization fn =
    let cache = new ConcurrentDictionary<_,_>()
    printfn "counter: %A" cache.Count

    fun arg ->
        printfn "%A is received" arg
        match cache.TryGetValue arg with
        | true, value -> value
        | false, _ ->
            let value = lazy fn arg
            cache.TryAdd(arg, value.Force()) |> ignore
            value.Force()

let shorter_memoization fn = 
    let cache = new ConcurrentDictionary<_,_>()
    fun arg -> cache.GetOrAdd(arg, lazy fn arg).Value

let memoization_2args fn = shorter_memoization (fun (arg1, arg2) -> fn arg1 arg2) |> fun fn' arg1 arg2 -> fn' (arg1, arg2)

let memoization_3args fn = shorter_memoization (fun (arg1, arg2, arg3) -> fn arg1 arg2 arg3) |> fun fn' arg1 arg2 arg3 -> fn' (arg1, arg2, arg3)