Memoization in Go

Memoization in Go

·

5 min read

Memoization (not to be confused with memorization) is an optimization technique used to cache a result of an expensive function call and return the cached result, when the function is called again with the same input parameters.

Memoization pattern is extensively used in functional programming world because it plays nicely with a concept of pure functions. What is a pure function?

A function is pure if it always returns the same output for the same input parameters and it has no side effects. In other words, it just maps input parameters to a result in a mathematical sense of it, and doesn't mutate any shared state along the way.

If we have a pure function that does an expensive computation for a provided input, it should be safe to cache a result of that computation when doing the computation for the first time for that input, and for subsequent calls just return the cached result without repeating the same computation again.

Apparently, it is trivial to achieve this caching behavior using a plain map structure as a cache and some imperative construct to save and retrieve values from the map. But is there a more functional and generic way to do it?

In functional programming it could be achieved by memoizing a function: an expensive function is wrapped into another function with the same signature but with the caching logic inside. The memoized function would cache results of function calls and would return back a cached result for the same input, if requested again. For example, it could look like that:

var memoF = Memo(func(a int) string {
    // expensive computation:
    time.Sleep(time.Millisecond * time.Duration(a))
    return fmt.Sprint(a)
})

r := memoF(2) // the first call is slow
r = memoF(2)  // other calls are fast

fmt.Println(r)
// Output: 2

But how would we implement the above Memo wrapper? Unsurprisingly, the implementation is quite straightforward:

func Memo[A comparable, B any](f func(A) B) func(A) B {
    m := make(map[A]B)

    return func(a A) B {
        result, ok := m[a]
        if ok {
            return result
        } else {
            evaluated := f(a)
            m[a] = evaluated
            return evaluated
        }
    }
}

As expected, this generic Memo function takes a 1-argument function as a parameter and returns a new function of exactly the same signature. This new function checks if the input has already been cached in the map, and returns it from the map if it's the case. Otherwise, it evaluates the original function and saves the result in the cache for future use.

One problem remains is that our Memo supports only 1-argument functions like func(A) B. But could we somehow reuse it to support functions of a different arity like func(A, B) C, func(A, B, C) D...?

One functional way to do it is to convert a function of 2 (3, or more) arguments into a function that takes 1 Tuple argument instead. A Tuple is a structure that contains a fixed number of elements, each with its own type, which could be defined like that:

type Tuple2[A, B any] struct {
    a A
    b B
}

type Tuple3[A, B, C any] struct {
    a A
    b B
    c C
}

func Tup2[A, B any](a A, b B) Tuple2[A, B] {
    return Tuple2[A, B]{a, b}
}

func Tup3[A, B, C any](a A, b B, c C) Tuple3[A, B, C] {
    return Tuple3[A, B, C]{a, b, c}
}

Then functions to convert a function of 2 (3, or more) arguments into a function that takes 1 Tuple argument could be defined like that:

func Tupled2[A, B, C any](f func(A, B) C) func(Tuple2[A, B]) C {
    return func(t Tuple2[A, B]) C {
        return f(t.a, t.b)
    }
}

func Tupled3[A, B, C, D any](f func(A, B, C) D) func(Tuple3[A, B, C]) D {
    return func(t Tuple3[A, B, C]) D {
        return f(t.a, t.b, t.c)
    }
}

And the opposite operations to convert a function with 1 tupled argument into a regular function of 2 (3, or more) arguments:

func UnTupled2[A, B, C any](f func(Tuple2[A, B]) C) func(A, B) C {
    return func(a A, b B) C {
        return f(Tup2(a, b))
    }
}

func UnTupled3[A, B, C, D any](f func(Tuple3[A, B, C]) D) func(A, B, C) D {
    return func(a A, b B, c C) D {
        return f(Tup3(a, b, c))
    }
}

Now using these Tupled/UnTupled functions we can make any function of 2 (3, or more) arguments compatible with our Memo function that expects a function of 1 argument:

func Memo2[A, B comparable, C any](f func(A, B) C) func(A, B) C {
    return UnTupled2(Memo(Tupled2(f)))
}

func Memo3[A, B, C comparable, D any](f func(A, B, C) D) func(A, B, C) D {
    return UnTupled3(Memo(Tupled3(f)))
}

And now we finally can memoize functions of 2 (3, or more) arguments also. For example, for a 2-argument function, usage of our Memo could look like that:

var memoF = Memo2(func(a, b int) string {
    // expensive computation:
    time.Sleep(time.Millisecond * time.Duration(a+b))
    return fmt.Sprint(a) + fmt.Sprint(b)
})

r := memoF(2, 2) // the first call is slow
r = memoF(2, 2)  // other calls are fast

fmt.Println(r)
// Output: 22

Conclusion

In this article we looked into Memoization technique, discussed its benefits, and a possible implementation of this pattern in Go. If you are interested to dive deeper into the functional programming topic, you are welcomed to join me on GitHub, where I develop a functional programming library for Go: github.com/ialekseev/go4fun. Thanks for reading!