How to perform thread-safe function memoization in c#?

Question

Here on stack overflow I've found the code that memoizes single-argument functions:

static Func<A, R> Memoize<A, R>(this Func<A, R> f)
{
    var d = new Dictionary<A, R>();
    return a=> 
    {
        R r;
        if (!d.TryGetValue(a, out r))
        {
            r = f(a);
            d.Add(a, r);
        }
        return r;
    };
}

While this code does its job for me, it fails sometimes when the memoized function is called from the multiple threads simultaneously: the Add method gets called twice with the same argument and throws an exception.

How can I make the memoization thread-safe?

Answer 1

You can use ConcurrentDictionary.GetOrAdd which does everything you need:

static Func<A, R> ThreadsafeMemoize<A, R>(this Func<A, R> f)
{
    var cache = new ConcurrentDictionary<A, R>();

    return argument => cache.GetOrAdd(argument, f);
}

The function f should be thread-safe itself, because it can be called from multiple threads simultaneously.

This code also doesn't guarantee that function f is called only once per unique argument value. It can be called many times, in fact, in the busy environment. If you need this kind of contract, you should take a look at the answers in this related question, but be warned that they're not as compact and require using locks.

Answer 2

Expanding on GMan's answer, I wanted to memoize a function with more than one argument. Here's how I did it, using a C# Tuple (requires C# 7) as they key for the ConcurrentDictionary.

This technique could easily be extended to allow yet more arguments:

public static class FunctionExtensions
{
    // Function with 1 argument
    public static Func<TArgument, TResult> Memoize<TArgument, TResult>
    (
        this Func<TArgument, TResult> func
    )
    {
        var cache = new ConcurrentDictionary<TArgument, TResult>();

        return argument => cache.GetOrAdd(argument, func);
    }

    // Function with 2 arguments
    public static Func<TArgument1, TArgument2, TResult> Memoize<TArgument1, TArgument2, TResult>
    (
        this Func<TArgument1, TArgument2, TResult> func
    )
    {
        var cache = new ConcurrentDictionary<(TArgument1, TArgument2), TResult>();

        return (argument1, argument2) =>
            cache.GetOrAdd((argument1, argument2), tuple => func(tuple.Item1, tuple.Item2));
    }
}

For example:

Func<int, string> example1Func = i => i.ToString();
var example1Memoized = example1Func.Memoize();
var example1Result = example1Memoized(66);

Func<int, int, int> example2Func = (a, b) => a + b;
var example2Memoized = example2Func.Memoize();
var example2Result = example2Memoized(3, 4);

(Of course, to get the benefit of memoization you'd normally want to keep example1Memoized / example2Memoized in a class variable or somewhere where they're not short-lived).