Tackling recursion (not just) in Typescript

Recursion allows for solving a certain domain of problems with clarity, conciseness and elegance. Sadly, using recursion in TypeScript (Javascript) comes at a price. In this post, we’ll show what specific issue can knock us down while embracing recursion, and how to overcome it.

The definition

Something is said to be recursive if defined in terms of itself. 

In TypeScript, recursion can be utilised for defining recursive types and recursive functions. In the code examples to follow, we’ll show both. Let’s create a type representing a general purpose tree structure first:

interface Tree<A> {
    value: A;
    subtrees: ReadonlyArray<Tree<A>>;
}

You can see that every node in our tree includes a value of type A, and can have an unbounded number of subtrees.

If we wanted to count the nodes our tree consists of using recursion, we could do it like this:

const countNodes = <A>(tree: Tree<A>): number =>
    1 + tree.subtrees.reduce(<A>(count: number, subtree: Tree<A>) => count + countNodes(subtree), 0);

This code is very concise and readable — basically we’ve defined that we get tree node count by adding 1 to the sum of the node count of all subtrees (we could even verify the correctness of our function using mathematical induction, but it’s beyond the scope of this post).

Everything feels fine and dandy with this piece of code in production, until we come across a deep enough tree and try to count its nodes.

By Heath Cajandig - https://www.flickr.com/photos/96228372@N06/26795406264/, CC BY 2.0, https://commons.wikimedia.org/w/index.php?curid=96920363

The problem

Here comes the distant cousin of an error so famous that it became the name of our precious lord and saviour - Stack Overflow. Fun fact: do not try to search for “stack overflow” on Stack Overflow. It’s not that the internet would go boom, it just simply won't find any relevant answers 😃.

In order to understand the reason for the error, we need to learn how function invocation works under the hood. Whenever a Javascript engine calls a function, it saves the function context to a place called the stack. As soon as function execution finishes, that context is taken off the stack. If we call a function within another function, the new context is stacked on top of the previous one. This context is very important. Among other things, it keeps the information about where to continue with the execution when we return from another function. As you can probably guess, with recursive function executed stack size can grow very quickly, and sooner or later, we hit the stack size limits defined in the Javascript engine 💥. 

The remedy

To avoid blowing the stack, we need to make our function tail recursive — meaning a recursive call is the last thing executed by the function. Unfortunately, this condition is not met in our countNodes function. The last thing we do is actually adding 1 to the sum of the node count of all subtrees. We can, however, easily transform function as follows:

const countNodesTailRecursive = <A>(tree: Tree<A>): number => {
    const iterate = <A>(subtrees: ReadonlyArray<Tree<A>>, count: number = 0): number =>
        subtrees.length
            ? iterate(subtrees.flatMap(subtree => subtree.subtrees), count + subtrees.length)
            : count;
    return iterate([tree], 0);
}

We introduced a local iterate function that does the computation. We carry the intermediate result here as the last argument, so the call to iterate can be the last thing executed. That's why there is actually no need for function context, precisely because we just return the value from the function.

If we tried to run the changed code, our stack would still go boom 😭. Most of the functional programming-oriented languages do tail call optimization (haskell, scala, f#) and do not push the new function context to the stack if that function is tail recursive. This is not the case with Javascript, as the majority of Javascript engines [do not support this feature](https://kangax.github.io/compat-table/es6# test-propertail_calls(tail_call_optimisation)) (to be fair, Node.js had tail call optimization till version 8 as an experimental flag, but they’ve dropped it 🤔). 

Nothing is lost though — we can optimise the code by ourselves and the change is rather trivial:

const countNodesTailRecursiveTrampolinedCustom = <A>(tree: Tree<A>): number => {
    const iterate = <A>(subtrees: ReadonlyArray<Tree<A>>, count: number = 0) =>
        subtrees.length
            ? () => iterate(subtrees.flatMap(subtree => subtree.subtrees), count + subtrees.length)
            : count;

    let result =  iterate([tree], 0)
    while (result instanceof Function) {
        result = result()
    }

    return result;
}

Instead of recursively calling iterate directly, we return the function that calls it. This way we prevent stacking contexts, and simply call the result in a loop until there is a final result.

This simple technique is often called trampolining, and no, I wont search for any funny picture with a trampoline in it 😉.

The generalisation

To prevent us from duplicating the code again and again, we can generalise the trampolining and type safety at the same time.

export type Trampoline<A> = Recurse<A> | Value<A>

export interface Recurse<A> {
    _tag: 'Recurse',
    recurse: () => Trampoline<A>,
}

export interface Value<A> {
    _tag: 'Value',
    value: A,
}


const recurse = <A>(f: () => Trampoline<A>): Recurse<A> =>
    ({ _tag: 'Recurse', recurse: f })


const value = <A>(a: A): Value<A> =>
    ({ _tag: 'Value', value: a })

So what did we do here? We’ve prepared a type representing the return value from our recursive function. It is a union type describing either final result (Value), or a yet-to-called function to get the final result (Recurse). 

If you are wondering what the purpose of _tag property is, check TypeScript: Documentation - Narrowing for more details. Basically it allows us to do exhaustive checking in a switch case for each possible variant of the Trampoline type.

We also implemented the trampolined combinator. If passed on the tail recursive function which returns the Trampoline type, it provides us with a result function that does the trampolining trick:

const trampolined = <T extends ReadonlyArray<unknown>, A>(f: (...t: T) => Trampoline<A>) => (...t: T): A => {
    let result = f(...t)
    while (true) {
        switch (result._tag) {
            case "Recurse":
                result = result.recurse()
                break
            case "Value":
                return result.value
            default:
                absurd(result)
        }
    }
}

If you are unsure about that type-fu in the trampolined function signature, it’s used to correctly preserve types of function f arguments. Check the following github issue if you are interested in learning more.

With a few finishing touches to the iterate function, we can finally apply a trampolined combinator and start to celebrate the outcome!

const countNodesTailRecursiveTrampolined = <A>(tree: Tree<A>): number => {
    const iterate = <A>(subtrees: ReadonlyArray<Tree<A>>, count: number = 0): Trampoline<number> =>
        subtrees.length
            ? recurse(() => iterate(subtrees.flatMap(subtree => subtree.subtrees), count + subtrees.length))
            : value(count);


    return trampolined(iterate)([tree]);
}

Wrapping it up

We have shown that, with those two rather simple transformations - making recursive function tail call recursive and trampolining, we can enjoy recursion even in TypeScript (Javascript) without worrying about our stack safety.

Stay safe and recursive!

Further reading:

• Implementing and Defunctionalizing Tail Calls in Javascript – UWTB
• Induction and Recursion

‍