Memoizing an nary function in TypeScript
 Published on
 Reading Time
 9 min read
 Words count
 1892 words
What is memoization?
When a function is likely to be called multiple times with the same input, it makes sense to store this input and the corresponding output and reuse it. We know that the function would return the same value every time for the same input. This act of storing results to make the function run faster is memoization.
In most problems, the problem itself might not be memoizable. We need to break the problem down to subproblems that are possible to memoize. Some of the best examples to understand memoization are the Fibonacci generator or the Factorial number generator. In these problems it’s very easy to see how the problem can be broken down into repetitive subproblems where the inputs are the same.
There are many great articles and posts that discuss these in detail. For the scope of this blog post, let’s proceed ahead with making the problem a bit more complex.
nary
function
Arity of a function is the number of arguments a function takes as input. nary
means the function takes n
arguments as input. Most memoization libraries, for example, lodash.memoize
, use the memoization technique for a 1ary or unary function (a function that takes only one input).
The problem
Sometimes, we have a function that takes 2 inputs, and we want to memoize that function. How do we approach this problem? What does memoization mean for a function with two inputs? If we look at a simple implementation of memoization, we have a map where we store our results.
function memoize<Input, Result>(fn: (input: Input) => Result) {
const memoMap = new Map<Input, Result>();
return function (input: Input): Result {
if (memoMap.has(input)) return memoMap.get(input)!;
const result = fn(input);
memoMap.set(input, result);
return result;
};
}
In this example, if we were to use two inputs, how would we store the result? Would we use a map of map or is it for a pair of inputs?
Just wrap the inputs in an array or object
Let’s consider combining the inputs into an array or an object to make it a single input. As JavaScript developers, we immediately know why creating that wrapper object or array would not work as the key of a map. They are passed by reference. Every time the function executes, a {}
or []
syntax creates a new object, and the map.has(obj)
check will fail.
To check the memo works, we can introduce a side effect in the original function.
function add(a, b) {
console.log("perform expensive computation");
return a + b;
}
const memoAdd = memoize(([a, b]) => add(a, b));
// first call saves the output
memoAdd([5, 10]);
// these calls must return result from the map
// but they won't because every [5,10] is a new reference
memoAdd([5, 10]);
memoAdd([5, 10]);
Every time we call memoAdd
with [5,10]
, it results in calling the original function. You can observe by running this and seeing the logs multiple times. This is because []
creates a new array and its reference becomes the key of the memoMap
. The same applies for objects {}
as well.
We can, of course, create the input once and use it multiple times —
const input = [5, 10];
memoAdd(input);
memoAdd(input);
It’s not realistic to assume that all our function calls share the same space in code. Usually, the function calls happen from different parts of the codebase, or it is harder to share the input reference.
Stringify
Can we stringify the input as string equality check is ‘by value’ instead of reference? As you might have already realized, this has its own drawbacks. There is also an extra parse step that must happen inside the function implementation. It can become counterproductive, as stringify and then parse will cost more computation in the process we are trying to optimize.
Also, the order of keys in an object are not certain. It takes the insertion order. So stringifying that requires special stringify methods. Even if it’s an array, the same problems occur if the nary input contains objects or arrays.
Why libraries use unary functions for memoization
As you can see, the methods we explored so far have severe drawbacks that the libraries cannot assume. It has to work for a general case. The general case is in the simplest form — a unary function.
Currying
One of the core issues is that we were trying so far to fit multiple arguments into a single argument. What if we treated each argument as a separate function. In this process, we are converting an nary function to multiple unary functions. It is known as ‘currying’.
In short, currying can be visualized as —
foo(a, b, c, d);
// becomes
curriedFoo(a)(b)(c)(d);
If we take the function add,
function add(a, b) {
return a + b;
}
currying this function would return the following function —
function curriedAdd(a) {
return (b) => a + b;
}
Now, we have a chain of nested unary functions. And we can memoize unary functions using any memoization library like lodash.memoize. Did we solve the problem of memoizing an nary function?
Nested maps
We still have to memoize them at each level. Otherwise, it would be incomplete. Let’s visualize the concept of what we achieve by currying —
// original function
function add(a, b) {
return a + b;
}
// curried function
function curriedAdd(a) {
return (b) => a + b;
}
// memoized and curried function
const memoizedCurriedAdd = memoize((a) => memoize((b) => a + b));
// calling memoized nary function
memoizedCurriedAdd(5)(10);
Memoize creates a map whose keys are inputs of the unary function and values are outputs. If you look at the memoizedCurriedAdd
definition, the result of the outer function is another memoized function. So we would essentially end up with the following nested map represented as a JS object —
map = {
[a]: memoizedFunction({
[b]: result,
}),
};
Another way to implement this memoized function is to simply store the results in a nested map ourselves — i.e., instead of using value as another memoized function, it will point to a map. But the concept remains the same.
Recreating the original nary function
A nice property of memoizing unary functions is that the original API is unchanged. The input and output types are the same.
type UnaryFn<Input, Result> = (input: Input) => Result;
function memoize<Input, Result>(
fn: UnaryFn<Input, Result>
): UnaryFn<Input, Result>;
What memoize returns has the exact same signature. Can we achieve this for our curried function? Can we uncurry
it? Of course, let’s take a look at how we can do it —
// memoized curried function
const memoizedCurriedAdd = memoize((a) => memoize((b) => a + b));
// memoized uncurried function
const memoizedAdd = (a, b) => memoizedCurriedAdd(a)(b);
Determining number of uncurries
The next step is to generalize uncurry function. We did not do this for the curry function because we have to also memoize every step of currying. But uncurry, it’s just a function call with a different API.
The JavaScript language supports variadic arguments in two forms —

Rest operator
function foo(...rest) { rest; // [list of arguments] }

arguments
variablefunction foo() { arguments; // { 0: first, 1: second, and so on.. } }
Because of variadic arguments support, we cannot reliably know from the function it’s arity. <function>.length
gives the arity if there are no variadic arguments. Otherwise, it’s just 0
. We need to know the arity of the function to do the number of uncurries.
So, our only option is to define uncurry for each arity. We also want the types to flow so that we use strict types everywhere. And that this entire process should not degrade the type info of the original function. Here is the uncurry function for various arities.
// uncurry into a 2 arity function
function uncurry2<A, B, R>(fn: (a: A) => (b: B) => R) {
return (a: A, b: B) => fn(a)(b);
}
// uncurry into a 3 arity function
function uncurry3<A, B, C, R>(fn: (a: A) => (b: B) => (c: C) => R) {
return (a: A, b: B, c: C) => fn(a)(b)(c);
}
// uncurry into a 4 arity function
function uncurry4<A, B, C, D, R>(
fn: (a: A) => (b: B) => (c: C) => (d: D) => R
) {
return (a: A, b: B, c: C, d: D) => fn(a)(b)(c)(d);
}
// you can extrapolate this if you need more
nary
memoize definition
The same problem that the uncurry faces, memoize faces too. Due to variadic arguments support, we cannot reliably define a single function for memoizing an nary function. So, we must define for each arity.
Using similar principles we need all the types to flow such that we get the type definition that is the same as the original function.
Let’s start by defining memoize2
that memoizes a 2arity function —
type Fn2<A, B, R> = (a: A, b: B) => R;
function memoize2<A, B, R>(fn: Fn2<A, B, R>): Fn2<A, B, R> {
return uncurry2(memoize((a: A) => memoize((b: B) => fn(a, b))));
}
Let’s reiterate what is happening in this function. We will unwrap it inside out.
memoize
is a unary memoization function. It takes a unary function and returns a unary function that’s memoized. The nested memoize does the currying and memoization together. It creates a nested memoization of
fn(a, b)
and returnsmemoizedCurriedFn(a)(b)
. uncurry2
removes the curry used for memoization and returns the original API of the functionfn
. As you can see in the input value and return value, the input function’s signature is
Fn2<A, B, R>
. The return value is also exactly the same.
We can extrapolate this and define memoize3
and memoize for higher arity functions.
type Fn3<A, B, C, R> = (a: A, b: B, c: C) => R;
function memoize3<A, B, C, R>(fn: Fn3<A, B, C, R>): Fn3<A, B, C, R> {
return uncurry3(
memoize((a: A) => memoize((b: B) => memoize((c: C) => fn(a, b, c))))
);
}
Example
That’s all folks! Let’s try out an example,
function add(a, b, c) {
return a + b + c;
}
const memoizedAdd = memoize3(add);
// would call the function
memoizedAdd(5, 10, 15);
// would return the memoized value
memoizedAdd(5, 10, 15);
Conclusion
It was an interesting journey trying to solve a simpletounderstand problem, but yet we went into various aspects of the language, type system, usage of generics for proper flow of types, and concepts like currying and uncurrying. I hope the post was useful in knowing a bit more about JavaScript or TypeScript.
If you have any doubts or comments or questions or fixes for this post, please feel free to tweet to me at @heisenbugger.