JavaScript memoization: Handy functional programming

Any time you're writing code, it means the code is going to be doing work for you. Many times, the burden of that work is transparent to us and our users. When it does make itself apparent, we have a performance issue on our hands. On the front-end this might manifest itself as lag or jank. On the back-end we may have higher latency or slow responses.

Sometimes there are ways of choosing faster algorithms or better data structures to resolve our performance issues, but not always. So what do we do when we’re too slow but can’t speed things up any more? One way of doing this in JavaScript is called memoization, which is a functional programming technique used to deal with costly pure functions.

With memoization, you “wrap” a function with another function that remembers the arguments and returns results of function calls. The result is essentially an in-memory cache that is transparent to your users. Let’s get a close look at memoization by implementing memoization around a simple but costly function.

A note on performance improvements: If you're seriously looking at improving performance, the first step is to measure current performance. Without clear measurements, you won’t have any way of identifying your bottlenecks or validating candidate solutions. For the sake of this post, I’ve written a small function in JavaScript that takes a function and its arguments and times the execution of that function. It looks like this:

const timeFn = (fn, args = []) => {
  const start = Date.now();
  const returnValue = fn(...args);
  const end = Date.now();
  console.log(`Execution took ${end - start} ms`);
  return returnValue;
}

The function we’ll use to test out memoization will add up numbers between zero and a number passed in as an argument. That function could be defined like this:

const summate = (countTo) => {
  let value = 0;
  for (let i = 0; i < countTo; i++) {
    value += i;
  }
  return value;
}

Using our timer function we get the following:

let sum = timeFn(summate, [1000000000]);
// sum is set to 499999999067109000
// The console logs: Execution took 1064 ms

Now that we have a costly function and a way of measuring that cost, let’s understand and implement memoization of that function.

Memoization is a way of wrapping a function and saving arguments and return values so that the next time an argument is passed in, we can check to see if we’ve done the work already. If we have already done the work, we can skip the heavy lifting and return the value we saved from before.

Memoization only works for pure functions, which are functions that return the same results when given the same arguments and that have no side effects. A function named add that takes two arguments and returns the sum would be a pure function, because every time you hand it 1 and 2 as arguments, you get 3 as a return value. By contrast, a function that takes a single argument and returns the sum of the argument and the current time in milliseconds is not a pure function. Any time you hand it 1, it will return a different value because the time is always different.

Our summate function above is a pure function, so it is a candidate for memoization. The steps we’re going to take to memoize a function look like this:

1. When a function is called, look at the arguments passed in.
2. If the arguments are present as a key in a lookup table (in our case, a JavaScript object), return the value of that lookup.
3. If not, run the function with the arguments.
4. Save the return value of the function to our lookup table with the arguments as the key.
5. Return the calculated value.

The above steps will ensure that any time our function is called with arguments we’ve seen before, we’ll return our previously calculated result and save some time. To accomplish this transparently to our users, we will create a higher-order function that takes a function as an argument and returns a new function that performs the steps above. Here is that higher-order function:

let simpleMemoizer = (fn) => {
  // Set up a lookup table within a closure so we have
  // access every time our inner function is called
  const lookupTable = {};
  return (arg) => {
    // Step 1
    if(arg in lookupTable) {
      // Step 2
      return lookupTable[arg];
    } else {
      // Step 3
      let returnValue = fn(arg);
      // Step 4
      lookupTable[arg] = returnValue;
      // Step 5
      return returnValue;
    }
  }
}

We can then use this function to wrap our summate function to memoize it:

const memoSummate = simpleMemoizer(summate);

The resulting function, memoSummate, can be used just like summate, but it will immediately return results that it has already calculated instead of recalculating them. We can verify this by running a few examples through timeFn:

timeFn(memoSummate, [100000000])
// Logs: Execution took 115 ms
// Returns: 4999999950000000
// When we call it again…
timeFn(memoSummate, [100000000])
// Logs: Execution took 0 ms
// Returns: 4999999950000000

And there you have it! Note that this implementation of a memoizer is very naive for demonstration purposes. It only pays attention to a single argument and that argument essentially has to be a string or a number. More sophisticated memoization functions, like those provided in the lodash or ramda GitHub repositories, will cover those extra cases for you and provide more universal functionality.

Comments or requests for other functional programming topics? Continue the conversation on Twitter: @freethejazz.