< Back to Latest Articles

Using memoization to speed up slow functions

In this article you’ll learn how memoization can dramatically boost the performance of slow functions, and how easy Swift makes it thanks to its generics and closures.

Memoization allows us to speed up function calls by caching the return value for a particular piece of input. The technique is startlingly simple in Swift, and can deliver huge performance boosts in all kinds of apps and games.

Tip: I’ve had folks email me trying correct my spelling here, so just to be clear it is called “memoization” and not “memorization” – they both come from the same root Latin word, memoro, meaning “to remember.”

• Download the project for this tutorial here.

Quick links

• What’s the problem?
• Manual memoization
• Automatic memoization
• Why the slow down?
• Recursive memoization
• What now?
• Challenges

What’s the problem?

Memoization lets us cache the values of slow functions. For example, if you were building a game and wanted to keep a high frame rate while also performing slow operations such as calculating square roots or sine values, memoization is perfect for the job.

In fact, memoization is used extensively in functional programming, because if you have a pure function – a function that always returns identical output for given input because it relies on no external state and has no side effects – then there’s not much of a reason why the result shouldn’t be cached.

The example project I’ve given you is trivial: it has a UITableViewController with 40 rows. Inside each row is being shown the Fibonacci number equivalent for that row’s index path, so if you run the app you’ll see 0, 1, 1, 2, 3, 5, 8, 13, and so on. Each sequence is calculated by starting with 0 and 1, then calculating each new number by adding the previous two.

You’ll see the code calculates the correct output for any given input like this:

func fibonacci(_ number: Int) -> Int {
    number < 2 ? number : fibonacci(number - 1) + fibonacci(number - 2)
}

So, for the 0th and 1st items we just send back 0 and 1 directly, but for everything else the function recursively calls itself to add the previous numbers together.

If you try using the app you’ll see it has a problem: scrolling gets worse the lower you go, and in fact if you go towards the bottom you’ll see your device completely freeze up. It doesn’t matter how powerful your device is, this is extremely slow and you just won’t get even close to smooth scrolling.

The problem is that our simple-looking function is recursive, so it generates huge amount of work in just one line of code. For example, if we want to calculate the 50th value in the Fibonacci sequence:

• We need to calculate the 49th value of the sequence, number - 1. That needs to calculate the 48th value, which needs to calculate the 47th value, and the 46th, the 45th, and so on.
• We also need to calculate the 48th value of the sequence, number - 2, which needs to do all the same work.

What actually happens is that end up repeating vast amounts of work, because every single number in the sequence tries to calculate every single number before it and the single number before that.

You can see the problem in action by changing the function to be more explicit:

func fibonacci(_ number: Int) -> Int {
    print("Calculating \(number)")
    return number < 2 ? number : fibonacci(number - 1) + fibonacci(number - 2)
}

When the app runs you’ll see a huge wall of print() statements – well past 10,000 and probably even past 20,000 depending on which device you’re using.

One way to address this problem is to calculate the sequence up front, then read values from that cache.

So, you could add a property to ViewController like this:

let numbers = Array(0...40).map(fibonacci)

And now you can change cellForRowAt to read values from that cache, like this:

cell.textLabel?.text = String(numbers[indexPath.row])

If you run the app again, you’ll see it now scrolls smoothly. However, our fast runtime performance has come at a cost: our app now launches significantly slower, because it’s doing all the calculation up front.

You might then try using a lazy map, like this:

let numbers = Array(0...40).lazy.map(fibonacci)

But unlike lazy properties, lazy maps don’t store their results – our app will now launch faster, but we’re back to the same terrible scrolling performance.

We can do better than this, and that’s where memoization comes in.

Before you continue, please revert your code back to use the fibonacci() function directly. That means removing the numbers property and using String(fibonacci(indexPath.row)) for the cell’s text label

Manual memoization

Our app has two problems: scrolling around is slow because we have to calculate complex values, and even after we calculate a value as soon as we scroll away that value gets discarded and needs to be recalculate every time. Memoization can help us solve both those problems.

Remember, memoization is the process of caching returned values from a function for particular in input. In its most raw state, that would mean adding a global dictionary variable to store input and outputs for our function. Place this outside the ViewController class, so it’s a global variable:

var fibonacciCache = [Int: Int]()

And now we can modify our fibonacci() function so that it starts by attempting to read the cache for the provided input. If we get a cache hit (we found the result for that input) then we send it back immediately, but if not we calculate it as before then store that value in the cache. So, the cache will start off empty, but fill up as the function runs repeatedly.

Here’s how that looks:

func fibonacci(_ number: Int) -> Int {
    if let value = fibonacciCache[number] {
        return value
    }

    let newValue = number < 2 ? number : fibonacci(number - 1) + fibonacci(number - 2)
    fibonacciCache[number] = newValue
    return newValue
}

If you run the app now you’ll see it works perfectly: it launches instantly, scrolling is smooth all the way to the bottom, and if you scroll up and down then it doesn’t need to recalculate values – it’s all cached.

Before you continue, please revert your code so that we don’t cache values. That means removing the fibonacciCache dictionary and putting the fibonacci() function back to this:

func fibonacci(_ number: Int) -> Int {
    number < 2 ? number : fibonacci(number - 1) + fibonacci(number - 2)
}

If you want, you can keep the caching fibonacci() function around but commented out – we’ll be looking at it again later.

Automatic memoization

Global functions are common, but global variables are usually a bad idea. What we really want is for that cache dictionary to be stored inside the function, so we don’t make a mess in our code. Even worse, if we have several things we want to memoize then we’ll create an even bigger mess!

To fix this, we’re going to create a new function called memoize() that will be able to memoize any function that accepts a single value and returns a single value.

This needs some careful thinking:

1. If we’re going to be able to memoize anything, then we can’t just have a single global variable any more – we need that storage be inside the memoizing function.
2. To memoize any kind of data we need to use generics: a generic input and a generic output.
3. We can’t memoize absolutely everything, because dictionaries can only store keys that conform to the Hashable protocol.

The reason this is all possible in Swift is because it supports closures – anonymous functions that we can pass around as data. In this instance, that means our memoize() function will accept some sort of function (the work to do, e.g. fibonacci(), then send back a new function that wraps it in a cache. Because that function parameter will be saved, we need to mark it as @escaping so Swift will keep it alive.

Add this function below fibonacci():

// work with any sort of input and output as long as the input is hashable, accept a function that takes Input and returns Output, and send back a function that accepts Input and returns Output
func memoize<Input: Hashable, Output>(_ function: @escaping (Input) -> Output) -> (Input) -> Output {
    // our item cache
    var storage = [Input: Output]()

    // send back a new closure that does our calculation
    return { input in
        if let cached = storage[input] {
            return cached
        }

        let result = function(input)
        storage[input] = result
        return result
    }
}

The magic happens there because we’re returning a closure that will capture the storage cache. Inside that closure we’re calling the function that was passed in (fibonacci()) and storing it in the cache.

When it comes to using that memoize() function, you first add it as a property to ViewController, like this:

let memoizedFibonacci = memoize(fibonacci)

And now update cellForRowAt to call that memoized version, like this:

cell.textLabel?.text = String(memoizedFibonacci(indexPath.row))

That’s it! Remember, our memoize() function accepts as its parameter some sort of function. Importantly, memoize() matches the signature as the function it accepts, so if the function accepts an Int and returns a String then so does `memoize() – that’s why we can use it directly with our table view cell.

If you run the app now you’ll see it kind of works: our scrolling is slow at first, but then when you scroll away and come back it’s fast. This means we’ve only solved half the problem: our values aren’t being recalculated, but we’re still stuck with them being slow the first time they are being calculated.

So, what happened? How come our awesome generic function was much slower than the specialized function?

Why the slow down?

First, let’s take a look at the specialized caching fibonacci() function we wrote earlier:

func fibonacci(_ number: Int) -> Int {
    if let value = fibonacciCache[number] {
        return value
    }

    let newValue = number < 2 ? number : fibonacci(number - 1) + fibonacci(number - 2)
    fibonacciCache[number] = newValue
    return newValue
}

That attempts to find a value in the cache, but if it fails it will call itself to calculate the value.

Now look at our generic function again:

func memoize<Input: Hashable, Output>(_ function: @escaping (Input) -> Output) -> (Input) -> Output {
    // our item cache
    var storage = [Input: Output]()

    // send back a new closure that does our calculation
    return { input in
        if let cached = storage[input] {
            return cached
        }

        let result = function(input)
        storage[input] = result
        return result
    }
}

That attempts to find a value in the cache, but if it fails it will call the original function to calculate the value.

The difference might seem small, but it accounts for the massive performance difference: in the specialized function the cache was inside the function, so when the function called itself it would benefit from all previous cached values; but in the generic function the cache was only used for the original input value, with all other values going through the uncached original function directly.

For small examples like the above this isn’t a problem, but if you wanted the 40th item in the sequence then you’re doing thousands of times more work than necessary – every item has to be calculated again and again because it isn’t in the cache. When it finishes, Swift will remember the return value of fibonacci(40) and make that one specific function call faster, but if you then asked for fibonacci(41) then it will do everything from scratch.

Fixing this is hard. I mean, really hard. If I explain it here and you think “Paul finds this easy,” you should know that I first saw this technique used six years ago at WWDC – I’ve had six years to come to terms with this!

And you know, a lot of time you don’t need the solution I’m about to present. If you’re doing slow work that you want to cache but that isn’t recursive – any function that needs to call itself, such as the Fibonacci sequence or factorials – then you can stop here and our simple memoize() function is enough.

But if you do have memoization needs that involve recursive functions, or if you have a morbid fascination for just how hard this is going to be, read on…

Recursive memoization

We need to update our memoize() function so that rather than calling its input function it inside calls the caching closure that gets returned.

Having both forms of memoize() is helpful, so rather than change your existing function I suggest you take a copy of it that we can modify, calling the new version recursiveMemoize().

First, we need to change the parameter that our copy accepts. Right now it’s (Input) -> Output, which means it will take something in and send back something – an Int coming in and an Int going out, for the Fibonacci sequence.

We need to change that so we accept a function that accepts some input and returns some output, as well as the input it’s working with. So, the parameter becomes this: (Input) -> Output, Input) -> Output.

In isolation that isn’t easy, but it gets even more complex when viewed in light of the whole function signature:

func recursiveMemoize<Input: Hashable, Output>(_ function: @escaping ((Input) -> Output, Input) -> Output) -> (Input) -> Output {

That means:

• recursiveMemoize() accepts a function as its parameter.
• That function itself accepts a function, which must be able to accept our input and send back our output.

So we now have a function recursiveMemoize() that accepts a function as its parameter, and that function accepts it own function as a parameter. See, I told you it was hard?

Now, when it comes to the code inside recursiveMemoize() we need to make a few changes.

You see, the closure we return needs to able to call itself, which means it needs to have a name.

So, change the closure to this:

let memo = { input in

Now add return memo to the end of the function.

Next, inside that closure we’re calling our calculation function like this:

let result = function(input)

That worked fine when our function was (Input) -> Output, but now the function is ((Input) -> Output, Input) -> Output – it returns the same output value ultimately, but it needs an (Input) -> Output function as its first parameter.

Fortunately, that’s exactly what our memoization function does, so we can use that:

let result = function(memo, input)

And now we have a new problem: we’re using the memo() closure inside its own definition – we haven’t even finished creating it before we’re using it. Swift really doesn’t like this, so we need to make a change: we need to define the closure before we create it, like this:

var memo: ((Input) -> Output)!

memo = { input in

And that’s it: our new memoization function is ready. Again, this is complex Swift code so don’t feel bad if you need to go over it many times to really understand how it works.

However, if you thought that was hard wait until you see how we use it.

First, change the memoizedFibonacci property so that it starts a closure that accepts two parameters, a function, and an input:

let memoizedFibonacci = recursiveMemoize { fibonacci, number in

Now copy in the body of the existing fibonacci() function, like this:

let memoizedFibonacci = recursiveMemoize { fibonacci, number in
    number < 2 ? number : fibonacci(number - 1) + fibonacci(number - 2)
}

Yes, that closure now calls as a function one of the parameters it gets passed in – it’s recursive! That means it’s no longer using the global fibonacci() function we made earlier, but is instead calling itself.

This seems like it ought to be impossible, but it really isn’t: eventually the function call will resolve in number` being less than 2 so it will stop calling itself.

If you run the app now you should see we get smooth scrolling. So, we finally solved our problem fully, while also making it generic so that we can use it with any other kind of similar function.

What now?

Our original, simpler generic memoization function is useful in many places, and it’s the kind of thing you’ll find yourself using in any place where you have a complex, repeatable transformation being applied.

Like I said at the beginning, you could memoize square root calculations:

let memoizedSqrt = memoize(sqrt)

You could also memorize calculating the sine of numbers, although because there are several different forms of sine() you should be clear which one you mean:

let memoizedSin: (Double) -> Double = memoize(sin)

Notice how I’m using the original function name, just prefixed with memoized – that makes it easy to understand at a glance what’s happening when you come across the code elsewhere.

Before I’m doing, I want to add two important notes to this technique.

First, memoization is a time-space tradeoff: we save some time in exchange for chewing up a little extra memory to store all our cached data. If you have lots of spare CPU time and little spare memory, memoization might not be for you.

Second, recursive memoization allows us to do some tasks extremely quickly, but it’s not magic – recursive functions can still cause a stack overflow, where one function calls itself so much that the computer can’t handle it any more.

With great power comes great responsibility! And speed. Responsibility and speed. And power. Responsibility, speed, and power!

And flexibility.

Oh, you get the point…

Challenges

I have one question for you to think about, and one coding challenge. The question is this: if we had an app that uses Core Image filters with a slider to control filter strength, why would that make a poor choice for memoization?

Have a think about that before I answer. In the meantime, here’s the coding challenge: write a memoized factorial function, then benchmark how much faster it is with and without recursive memoization. You should use CFAbsoluteTimeGetCurrent() to measure timings.

As for the Core Image question, the problem with caching images is that they are big – they’ll chew up a lot of memory, which means our dictionary cache will fill up extremely quickly. It’s best to use memoization for simple data, such as String, Int, and Double.

If you liked this, you'd love Hacking with Swift+…

Here's just a sample of the other tutorials, with each one coming as an article to read and as a 4K Ultra HD video.

Find out more and subscribe here

---